<commit_before>package hoverfly

import (
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"path"
	"regexp"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/SpectoLabs/hoverfly/core/matching"
	"github.com/SpectoLabs/hoverfly/core/models"
	"github.com/SpectoLabs/hoverfly/core/views"
	"net/http"
)

// Import is a function that based on input decides whether it is a local resource or whether
// it should fetch it from remote server. It then imports given payload into the database
// or returns an error
func (hf *Hoverfly) Import(uri string) error {

	// assuming file URI is URL:
	if isURL(uri) {
		log.WithFields(log.Fields{
			"isURL":      isURL(uri),
			"importFrom": uri,
		}).Info("URL")
		return hf.ImportFromURL(uri)
	}
	// assuming file URI is disk location
	ext := path.Ext(uri)
	if ext != ".json" {
		return fmt.Errorf("Failed to import payloads, only JSON files are acceppted. Given file: %s", uri)
	}
	// checking whether it exists
	exists, err := exists(uri)
	if err != nil {
		return fmt.Errorf("Failed to import payloads from %s. Got error: %s", uri, err.Error())
	}
	if exists {
		// file is JSON and it exist
		return hf.ImportFromDisk(uri)
	}
	return fmt.Errorf("Failed to import payloads, given file '%s' does not exist", uri)
}

// URL is regexp to match http urls
const URL string = `^((ftp|https?):\/\/)(\S+(:\S*)?@)?((([1-9]\d?|1\d\d|2[01]\d|22[0-3])(\.(1?\d{1,2}|2[0-4]\d|25[0-5])){2}(?:\.([0-9]\d?|1\d\d|2[0-4]\d|25[0-4]))|(([a-zA-Z0-9]+([-\.][a-zA-Z0-9]+)*)|((www\.)?))?(([a-z\x{00a1}-\x{ffff}0-9]+-?-?)*[a-z\x{00a1}-\x{ffff}0-9]+)(?:\.([a-z\x{00a1}-\x{ffff}]{2,}))?))(:(\d{1,5}))?((\/|\?|#)[^\s]*)?$`

var rxURL = regexp.MustCompile(URL)

func isURL(str string) bool {
	if str == "" || len(str) >= 2083 || len(str) <= 3 || strings.HasPrefix(str, ".") {
		return false
	}
	u, err := url.Parse(str)
	if err != nil {
		return false
	}
	if strings.HasPrefix(u.Host, ".") {
		return false
	}
	if u.Host == "" && (u.Path != "" && !strings.Contains(u.Path, ".")) {
		return false
	}

	return rxURL.MatchString(str)

}

func exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return true, err
}

// ImportFromDisk - takes one string value and tries to open a file, then parse it into recordedRequests structure
// (which is default format in which Hoverfly exports captured requests) and imports those requests into the database
func (hf *Hoverfly) ImportFromDisk(path string) error {
	pairsFile, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("Got error while opening payloads file, error %s", err.Error())
	}

	var requests views.RequestResponsePairPayload

	jsonParser := json.NewDecoder(pairsFile)
	if err = jsonParser.Decode(&requests); err != nil {
		return fmt.Errorf("Got error while parsing payloads file, error %s", err.Error())
	}

	return hf.ImportRequestResponsePairViews(requests.Data)
}

// ImportFromURL - takes one string value and tries connect to a remote server, then parse response body into
// recordedRequests structure (which is default format in which Hoverfly exports captured requests) and
// imports those requests into the database
func (hf *Hoverfly) ImportFromURL(url string) error {

	resp, err := hf.HTTP.Get(url)
	if err != nil {
		return fmt.Errorf("Failed to fetch given URL, error %s", err.Error())
	}

	var requests views.RequestResponsePairPayload

	jsonParser := json.NewDecoder(resp.Body)
	if err = jsonParser.Decode(&requests); err != nil {
		return fmt.Errorf("Got error while parsing payloads, error %s", err.Error())
	}

	return hf.ImportRequestResponsePairViews(requests.Data)
}

func isJSON(s string) bool {
	var js map[string]interface{}
	return json.Unmarshal([]byte(s), &js) == nil

}

// ImportRequestResponsePairViews - a function to save given pairs into the database.
func (hf *Hoverfly) ImportRequestResponsePairViews(pairViews []views.RequestResponsePairView) error {
	if len(pairViews) > 0 {
		success := 0
		failed := 0
		for _, pairView := range pairViews {
			if pairView.Request.RequestType == "template" {
				responseDetails := models.NewResponseDetialsFromResponseDetailsView(pairView.Response)

				getValue := func (value string) *string {
					if value == "" {
						return nil
					}
					return &value
				}

				requestTemplate := matching.RequestTemplate{
					Path: getValue(pairView.Request.Path),
					Method: getValue(pairView.Request.Method),
					Destination: getValue(pairView.Request.Destination),
					Scheme: getValue(pairView.Request.Scheme),
					Query: getValue(pairView.Request.Query),
					Body: getValue(pairView.Request.Body),
					Headers: pairView.Request.Headers,
				}

				requestTemplateResponsePair := matching.RequestTemplateResponsePair{
					RequestTemplate: requestTemplate,
					Response: responseDetails,
				}

				hf.RequestMatcher.TemplateStore = append(hf.RequestMatcher.TemplateStore, requestTemplateResponsePair)
				continue
			}

			// Convert PayloadView back to Payload for internal storage
			pair := models.NewRequestResponsePairFromRequestResponsePairView(pairView)

			if len(pair.Request.Headers) == 0 {
				pair.Request.Headers = make(map[string][]string)
			}

			if _, present := pair.Request.Headers["Content-Type"]; !present {
				// sniffing content types
				if isJSON(pair.Request.Body) {
					pair.Request.Headers["Content-Type"] = []string{"application/json"}
				} else {
					ct := http.DetectContentType([]byte(pair.Request.Body))
					pair.Request.Headers["Content-Type"] = []string{ct}
				}
			}

			pairBytes, err := pair.Encode()
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
				}).Error("Failed to encode payload")
				failed++
			} else {
				// hook
				var en Entry
				en.ActionType = ActionTypeRequestCaptured
				en.Message = "imported"
				en.Time = time.Now()
				en.Data = pairBytes

				if err := hf.Hooks.Fire(ActionTypeRequestCaptured, &en); err != nil {
					log.WithFields(log.Fields{
						"error":      err.Error(),
						"message":    en.Message,
						"actionType": ActionTypeRequestCaptured,
					}).Error("failed to fire hook")
				}

				err := hf.RequestMatcher.SaveRequestResponsePair(&pair)
				if err != nil {
					log.WithFields(log.Fields{
						"error": err.Error(),
					}).Error("Failed to save payload")
				}

				if err == nil {
					success++
				} else {
					failed++
				}
			}
		}
		log.WithFields(log.Fields{
			"total":      len(pairViews),
			"successful": success,
			"failed":     failed,
		}).Info("payloads imported")
		return nil
	}
	return fmt.Errorf("Bad request. Nothing to import!")
}
<commit_msg>Added the success increment so that the success count is correct on the logs<commit_after>package hoverfly

import (
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"path"
	"regexp"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/SpectoLabs/hoverfly/core/matching"
	"github.com/SpectoLabs/hoverfly/core/models"
	"github.com/SpectoLabs/hoverfly/core/views"
	"net/http"
)

// Import is a function that based on input decides whether it is a local resource or whether
// it should fetch it from remote server. It then imports given payload into the database
// or returns an error
func (hf *Hoverfly) Import(uri string) error {

	// assuming file URI is URL:
	if isURL(uri) {
		log.WithFields(log.Fields{
			"isURL":      isURL(uri),
			"importFrom": uri,
		}).Info("URL")
		return hf.ImportFromURL(uri)
	}
	// assuming file URI is disk location
	ext := path.Ext(uri)
	if ext != ".json" {
		return fmt.Errorf("Failed to import payloads, only JSON files are acceppted. Given file: %s", uri)
	}
	// checking whether it exists
	exists, err := exists(uri)
	if err != nil {
		return fmt.Errorf("Failed to import payloads from %s. Got error: %s", uri, err.Error())
	}
	if exists {
		// file is JSON and it exist
		return hf.ImportFromDisk(uri)
	}
	return fmt.Errorf("Failed to import payloads, given file '%s' does not exist", uri)
}

// URL is regexp to match http urls
const URL string = `^((ftp|https?):\/\/)(\S+(:\S*)?@)?((([1-9]\d?|1\d\d|2[01]\d|22[0-3])(\.(1?\d{1,2}|2[0-4]\d|25[0-5])){2}(?:\.([0-9]\d?|1\d\d|2[0-4]\d|25[0-4]))|(([a-zA-Z0-9]+([-\.][a-zA-Z0-9]+)*)|((www\.)?))?(([a-z\x{00a1}-\x{ffff}0-9]+-?-?)*[a-z\x{00a1}-\x{ffff}0-9]+)(?:\.([a-z\x{00a1}-\x{ffff}]{2,}))?))(:(\d{1,5}))?((\/|\?|#)[^\s]*)?$`

var rxURL = regexp.MustCompile(URL)

func isURL(str string) bool {
	if str == "" || len(str) >= 2083 || len(str) <= 3 || strings.HasPrefix(str, ".") {
		return false
	}
	u, err := url.Parse(str)
	if err != nil {
		return false
	}
	if strings.HasPrefix(u.Host, ".") {
		return false
	}
	if u.Host == "" && (u.Path != "" && !strings.Contains(u.Path, ".")) {
		return false
	}

	return rxURL.MatchString(str)

}

func exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return true, err
}

// ImportFromDisk - takes one string value and tries to open a file, then parse it into recordedRequests structure
// (which is default format in which Hoverfly exports captured requests) and imports those requests into the database
func (hf *Hoverfly) ImportFromDisk(path string) error {
	pairsFile, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("Got error while opening payloads file, error %s", err.Error())
	}

	var requests views.RequestResponsePairPayload

	jsonParser := json.NewDecoder(pairsFile)
	if err = jsonParser.Decode(&requests); err != nil {
		return fmt.Errorf("Got error while parsing payloads file, error %s", err.Error())
	}

	return hf.ImportRequestResponsePairViews(requests.Data)
}

// ImportFromURL - takes one string value and tries connect to a remote server, then parse response body into
// recordedRequests structure (which is default format in which Hoverfly exports captured requests) and
// imports those requests into the database
func (hf *Hoverfly) ImportFromURL(url string) error {

	resp, err := hf.HTTP.Get(url)
	if err != nil {
		return fmt.Errorf("Failed to fetch given URL, error %s", err.Error())
	}

	var requests views.RequestResponsePairPayload

	jsonParser := json.NewDecoder(resp.Body)
	if err = jsonParser.Decode(&requests); err != nil {
		return fmt.Errorf("Got error while parsing payloads, error %s", err.Error())
	}

	return hf.ImportRequestResponsePairViews(requests.Data)
}

func isJSON(s string) bool {
	var js map[string]interface{}
	return json.Unmarshal([]byte(s), &js) == nil

}

// ImportRequestResponsePairViews - a function to save given pairs into the database.
func (hf *Hoverfly) ImportRequestResponsePairViews(pairViews []views.RequestResponsePairView) error {
	if len(pairViews) > 0 {
		success := 0
		failed := 0
		for _, pairView := range pairViews {
			if pairView.Request.RequestType == "template" {
				responseDetails := models.NewResponseDetialsFromResponseDetailsView(pairView.Response)

				getValue := func (value string) *string {
					if value == "" {
						return nil
					}
					return &value
				}

				requestTemplate := matching.RequestTemplate{
					Path: getValue(pairView.Request.Path),
					Method: getValue(pairView.Request.Method),
					Destination: getValue(pairView.Request.Destination),
					Scheme: getValue(pairView.Request.Scheme),
					Query: getValue(pairView.Request.Query),
					Body: getValue(pairView.Request.Body),
					Headers: pairView.Request.Headers,
				}

				requestTemplateResponsePair := matching.RequestTemplateResponsePair{
					RequestTemplate: requestTemplate,
					Response: responseDetails,
				}

				hf.RequestMatcher.TemplateStore = append(hf.RequestMatcher.TemplateStore, requestTemplateResponsePair)
				success++
				continue
			}

			// Convert PayloadView back to Payload for internal storage
			pair := models.NewRequestResponsePairFromRequestResponsePairView(pairView)

			if len(pair.Request.Headers) == 0 {
				pair.Request.Headers = make(map[string][]string)
			}

			if _, present := pair.Request.Headers["Content-Type"]; !present {
				// sniffing content types
				if isJSON(pair.Request.Body) {
					pair.Request.Headers["Content-Type"] = []string{"application/json"}
				} else {
					ct := http.DetectContentType([]byte(pair.Request.Body))
					pair.Request.Headers["Content-Type"] = []string{ct}
				}
			}

			pairBytes, err := pair.Encode()
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
				}).Error("Failed to encode payload")
				failed++
			} else {
				// hook
				var en Entry
				en.ActionType = ActionTypeRequestCaptured
				en.Message = "imported"
				en.Time = time.Now()
				en.Data = pairBytes

				if err := hf.Hooks.Fire(ActionTypeRequestCaptured, &en); err != nil {
					log.WithFields(log.Fields{
						"error":      err.Error(),
						"message":    en.Message,
						"actionType": ActionTypeRequestCaptured,
					}).Error("failed to fire hook")
				}

				err := hf.RequestMatcher.SaveRequestResponsePair(&pair)
				if err != nil {
					log.WithFields(log.Fields{
						"error": err.Error(),
					}).Error("Failed to save payload")
				}

				if err == nil {
					success++
				} else {
					failed++
				}
			}
		}
		log.WithFields(log.Fields{
			"total":      len(pairViews),
			"successful": success,
			"failed":     failed,
		}).Info("payloads imported")
		return nil
	}
	return fmt.Errorf("Bad request. Nothing to import!")
}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"

	"github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
const Version = "0.9.6"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
var VersionPrerelease string

// SemVersion is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVersion = version.Must(version.NewVersion(Version))

// VersionHeader is the header name used to send the current terraform version
// in http requests.
const VersionHeader = "Terraform-Version"

func VersionString() string {
	if VersionPrerelease != "" {
		return fmt.Sprintf("%s-%s", Version, VersionPrerelease)
	}
	return Version
}
<commit_msg>core: use codified default for prerelease string<commit_after>package terraform

import (
	"fmt"

	"github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
const Version = "0.9.6"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
var VersionPrerelease = "dev"

// SemVersion is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVersion = version.Must(version.NewVersion(Version))

// VersionHeader is the header name used to send the current terraform version
// in http requests.
const VersionHeader = "Terraform-Version"

func VersionString() string {
	if VersionPrerelease != "" {
		return fmt.Sprintf("%s-%s", Version, VersionPrerelease)
	}
	return Version
}
<|endoftext|>


<commit_before>//
// DISCLAIMER
//
// Copyright 2017 ArangoDB GmbH, Cologne, Germany
//
// 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.
//
// Copyright holder is ArangoDB GmbH, Cologne, Germany
//
// Author Ewout Prangsma
//

package test

import (
	"context"
	"testing"
	"time"

	driver "github.com/arangodb/go-driver"
)

// TestClusterHealth tests the Cluster.Health method.
func TestClusterHealth(t *testing.T) {
	ctx := context.Background()
	c := createClientFromEnv(t, true)
	cl, err := c.Cluster(ctx)
	if driver.IsPreconditionFailed(err) {
		t.Skip("Not a cluster")
	} else {
		h, err := cl.Health(ctx)
		if err != nil {
			t.Fatalf("Health failed: %s", describe(err))
		}
		if h.ID == "" {
			t.Error("Expected cluster ID to be non-empty")
		}
		agents := 0
		dbservers := 0
		coordinators := 0
		for _, sh := range h.Health {
			switch sh.Role {
			case driver.ServerRoleAgent:
				agents++
			case driver.ServerRoleDBServer:
				dbservers++
			case driver.ServerRoleCoordinator:
				coordinators++
			}
		}
		if agents == 0 {
			t.Error("Expected at least 1 agent")
		}
		if dbservers == 0 {
			t.Error("Expected at least 1 dbserver")
		}
		if coordinators == 0 {
			t.Error("Expected at least 1 coordinator")
		}
	}
}

// TestClusterDatabaseInventory tests the Cluster.DatabaseInventory method.
func TestClusterDatabaseInventory(t *testing.T) {
	ctx := context.Background()
	c := createClientFromEnv(t, true)
	cl, err := c.Cluster(ctx)
	if driver.IsPreconditionFailed(err) {
		t.Skip("Not a cluster")
	} else {
		db, err := c.Database(ctx, "_system")
		if err != nil {
			t.Fatalf("Failed to open _system database: %s", describe(err))
		}
		h, err := cl.Health(ctx)
		if err != nil {
			t.Fatalf("Health failed: %s", describe(err))
		}
		inv, err := cl.DatabaseInventory(ctx, db)
		if err != nil {
			t.Fatalf("DatabaseInventory failed: %s", describe(err))
		}
		if len(inv.Collections) == 0 {
			t.Error("Expected multiple collections, got 0")
		}
		for _, col := range inv.Collections {
			if len(col.Parameters.Shards) == 0 {
				t.Errorf("Expected 1 or more shards in collection %s, got 0", col.Parameters.Name)
			}
			for shardID, dbServers := range col.Parameters.Shards {
				for _, serverID := range dbServers {
					if _, found := h.Health[serverID]; !found {
						t.Errorf("Unexpected dbserver ID for shard '%s': %s", shardID, serverID)
					}
				}
			}
		}
	}
}

// TestClusterMoveShard tests the Cluster.MoveShard method.
func TestClusterMoveShard(t *testing.T) {
	ctx := context.Background()
	c := createClientFromEnv(t, true)
	cl, err := c.Cluster(ctx)
	if driver.IsPreconditionFailed(err) {
		t.Skip("Not a cluster")
	} else {
		db, err := c.Database(ctx, "_system")
		if err != nil {
			t.Fatalf("Failed to open _system database: %s", describe(err))
		}
		col, err := db.CreateCollection(ctx, "test_move_shard", &driver.CreateCollectionOptions{
			NumberOfShards: 12,
		})
		if err != nil {
			t.Fatalf("CreateCollection failed: %s", describe(err))
		}
		h, err := cl.Health(ctx)
		if err != nil {
			t.Fatalf("Health failed: %s", describe(err))
		}
		inv, err := cl.DatabaseInventory(ctx, db)
		if err != nil {
			t.Fatalf("DatabaseInventory failed: %s", describe(err))
		}
		if len(inv.Collections) == 0 {
			t.Error("Expected multiple collections, got 0")
		}
		var targetServerID driver.ServerID
		for id, s := range h.Health {
			if s.Role == driver.ServerRoleDBServer {
				targetServerID = id
				break
			}
		}
		if len(targetServerID) == 0 {
			t.Fatalf("Failed to find any dbserver")
		}
		movedShards := 0
		for _, colInv := range inv.Collections {
			if colInv.Parameters.Name == col.Name() {
				for shardID, dbServers := range colInv.Parameters.Shards {
					if dbServers[0] != targetServerID {
						movedShards++
						if err := cl.MoveShard(ctx, col, shardID, dbServers[0], targetServerID); err != nil {
							t.Errorf("MoveShard for shard %s in collection %s failed: %s", shardID, col.Name(), describe(err))
						}
					}
				}
			}
		}
		if movedShards == 0 {
			t.Fatal("Expected to have moved at least 1 shard, all seem to be on target server already")
		}
		// Wait until all shards are on the targetServerID
		start := time.Now()
		maxTestTime := time.Minute
		lastShardsNotOnTargetServerID := movedShards
		for {
			shardsNotOnTargetServerID := 0
			inv, err := cl.DatabaseInventory(ctx, db)
			if err != nil {
				t.Errorf("DatabaseInventory failed: %s", describe(err))
			} else {
				for _, colInv := range inv.Collections {
					if colInv.Parameters.Name == col.Name() {
						for shardID, dbServers := range colInv.Parameters.Shards {
							if dbServers[0] != targetServerID {
								shardsNotOnTargetServerID++
								t.Logf("Shard %s in on %s, wanted %s", shardID, dbServers[0], targetServerID)
							}
						}
					}
				}
			}
			if shardsNotOnTargetServerID == 0 {
				// We're done
				break
			}
			if shardsNotOnTargetServerID != lastShardsNotOnTargetServerID {
				// Something changed, we give a bit more time
				maxTestTime = maxTestTime + time.Second*15
				lastShardsNotOnTargetServerID = shardsNotOnTargetServerID
			}
			if time.Since(start) > maxTestTime {
				t.Errorf("%d shards did not move within %s", shardsNotOnTargetServerID, maxTestTime)
				break
			}
			t.Log("Waiting a bit")
			time.Sleep(time.Second * 5)
		}
	}
}
<commit_msg>Include raw response<commit_after>//
// DISCLAIMER
//
// Copyright 2017 ArangoDB GmbH, Cologne, Germany
//
// 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.
//
// Copyright holder is ArangoDB GmbH, Cologne, Germany
//
// Author Ewout Prangsma
//

package test

import (
	"context"
	"testing"
	"time"

	driver "github.com/arangodb/go-driver"
)

// TestClusterHealth tests the Cluster.Health method.
func TestClusterHealth(t *testing.T) {
	ctx := context.Background()
	c := createClientFromEnv(t, true)
	cl, err := c.Cluster(ctx)
	if driver.IsPreconditionFailed(err) {
		t.Skip("Not a cluster")
	} else {
		h, err := cl.Health(ctx)
		if err != nil {
			t.Fatalf("Health failed: %s", describe(err))
		}
		if h.ID == "" {
			t.Error("Expected cluster ID to be non-empty")
		}
		agents := 0
		dbservers := 0
		coordinators := 0
		for _, sh := range h.Health {
			switch sh.Role {
			case driver.ServerRoleAgent:
				agents++
			case driver.ServerRoleDBServer:
				dbservers++
			case driver.ServerRoleCoordinator:
				coordinators++
			}
		}
		if agents == 0 {
			t.Error("Expected at least 1 agent")
		}
		if dbservers == 0 {
			t.Error("Expected at least 1 dbserver")
		}
		if coordinators == 0 {
			t.Error("Expected at least 1 coordinator")
		}
	}
}

// TestClusterDatabaseInventory tests the Cluster.DatabaseInventory method.
func TestClusterDatabaseInventory(t *testing.T) {
	ctx := context.Background()
	c := createClientFromEnv(t, true)
	cl, err := c.Cluster(ctx)
	if driver.IsPreconditionFailed(err) {
		t.Skip("Not a cluster")
	} else {
		db, err := c.Database(ctx, "_system")
		if err != nil {
			t.Fatalf("Failed to open _system database: %s", describe(err))
		}
		h, err := cl.Health(ctx)
		if err != nil {
			t.Fatalf("Health failed: %s", describe(err))
		}
		inv, err := cl.DatabaseInventory(ctx, db)
		if err != nil {
			t.Fatalf("DatabaseInventory failed: %s", describe(err))
		}
		if len(inv.Collections) == 0 {
			t.Error("Expected multiple collections, got 0")
		}
		for _, col := range inv.Collections {
			if len(col.Parameters.Shards) == 0 {
				t.Errorf("Expected 1 or more shards in collection %s, got 0", col.Parameters.Name)
			}
			for shardID, dbServers := range col.Parameters.Shards {
				for _, serverID := range dbServers {
					if _, found := h.Health[serverID]; !found {
						t.Errorf("Unexpected dbserver ID for shard '%s': %s", shardID, serverID)
					}
				}
			}
		}
	}
}

// TestClusterMoveShard tests the Cluster.MoveShard method.
func TestClusterMoveShard(t *testing.T) {
	ctx := context.Background()
	c := createClientFromEnv(t, true)
	cl, err := c.Cluster(ctx)
	if driver.IsPreconditionFailed(err) {
		t.Skip("Not a cluster")
	} else {
		db, err := c.Database(ctx, "_system")
		if err != nil {
			t.Fatalf("Failed to open _system database: %s", describe(err))
		}
		col, err := db.CreateCollection(ctx, "test_move_shard", &driver.CreateCollectionOptions{
			NumberOfShards: 12,
		})
		if err != nil {
			t.Fatalf("CreateCollection failed: %s", describe(err))
		}
		h, err := cl.Health(ctx)
		if err != nil {
			t.Fatalf("Health failed: %s", describe(err))
		}
		inv, err := cl.DatabaseInventory(ctx, db)
		if err != nil {
			t.Fatalf("DatabaseInventory failed: %s", describe(err))
		}
		if len(inv.Collections) == 0 {
			t.Error("Expected multiple collections, got 0")
		}
		var targetServerID driver.ServerID
		for id, s := range h.Health {
			if s.Role == driver.ServerRoleDBServer {
				targetServerID = id
				break
			}
		}
		if len(targetServerID) == 0 {
			t.Fatalf("Failed to find any dbserver")
		}
		movedShards := 0
		for _, colInv := range inv.Collections {
			if colInv.Parameters.Name == col.Name() {
				for shardID, dbServers := range colInv.Parameters.Shards {
					if dbServers[0] != targetServerID {
						movedShards++
						var rawResponse []byte
						if err := cl.MoveShard(driver.WithRawResponse(ctx, &rawResponse), col, shardID, dbServers[0], targetServerID); err != nil {
							t.Errorf("MoveShard for shard %s in collection %s failed: %s (raw response '%s')", shardID, col.Name(), describe(err), string(rawResponse))
						}
					}
				}
			}
		}
		if movedShards == 0 {
			t.Fatal("Expected to have moved at least 1 shard, all seem to be on target server already")
		}
		// Wait until all shards are on the targetServerID
		start := time.Now()
		maxTestTime := time.Minute
		lastShardsNotOnTargetServerID := movedShards
		for {
			shardsNotOnTargetServerID := 0
			inv, err := cl.DatabaseInventory(ctx, db)
			if err != nil {
				t.Errorf("DatabaseInventory failed: %s", describe(err))
			} else {
				for _, colInv := range inv.Collections {
					if colInv.Parameters.Name == col.Name() {
						for shardID, dbServers := range colInv.Parameters.Shards {
							if dbServers[0] != targetServerID {
								shardsNotOnTargetServerID++
								t.Logf("Shard %s in on %s, wanted %s", shardID, dbServers[0], targetServerID)
							}
						}
					}
				}
			}
			if shardsNotOnTargetServerID == 0 {
				// We're done
				break
			}
			if shardsNotOnTargetServerID != lastShardsNotOnTargetServerID {
				// Something changed, we give a bit more time
				maxTestTime = maxTestTime + time.Second*15
				lastShardsNotOnTargetServerID = shardsNotOnTargetServerID
			}
			if time.Since(start) > maxTestTime {
				t.Errorf("%d shards did not move within %s", shardsNotOnTargetServerID, maxTestTime)
				break
			}
			t.Log("Waiting a bit")
			time.Sleep(time.Second * 5)
		}
	}
}
<|endoftext|>


<commit_before>// +build !nacl
// run

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

// Run the sinit test.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
)

func cleanup() {
	os.Remove("linkmain.o")
	os.Remove("linkmain.a")
	os.Remove("linkmain1.o")
	os.Remove("linkmain1.a")
	os.Remove("linkmain.exe")
}

func run(cmdline string) {
	args := strings.Fields(cmdline)
	cmd := exec.Command(args[0], args[1:]...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		fmt.Printf("$ %s\n", strings.Join(args, " "))
		fmt.Println(string(out))
		fmt.Println(err)
		cleanup()
		os.Exit(1)
	}
}

func runFail(args ...string) {
	cmd := exec.Command(args[0], args[1:]...)
	out, err := cmd.CombinedOutput()
	if err == nil {
		fmt.Printf("$ %s\n", strings.Join(args, " "))
		fmt.Println(string(out))
		fmt.Println("SHOULD HAVE FAILED!")
		cleanup()
		os.Exit(1)
	}
}

func main() {
	// helloworld.go is package main
	run("go tool compile -o linkmain.o helloworld.go")
	run("go tool compile -pack -o linkmain.a helloworld.go")
	run("go tool link -o linkmain.exe linkmain.o")
	run("go tool link -o linkmain.exe linkmain.a")

	// linkmain.go is not
	run("go tool compile -o linkmain.o linkmain.go")
	run("go tool compile -pack -o linkmain.a linkmain.go")
	runFail("go tool link -o linkmain.exe linkmain1.o")
	runFail("go tool link -o linkmain.exe linkmain1.a")
	cleanup()
}
<commit_msg>test: fix linkmain test<commit_after>// +build !nacl
// run

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

// Run the sinit test.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
)

func cleanup() {
	os.Remove("linkmain.o")
	os.Remove("linkmain.a")
	os.Remove("linkmain1.o")
	os.Remove("linkmain1.a")
	os.Remove("linkmain.exe")
}

func run(cmdline string) {
	args := strings.Fields(cmdline)
	cmd := exec.Command(args[0], args[1:]...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		fmt.Printf("$ %s\n", cmdline)
		fmt.Println(string(out))
		fmt.Println(err)
		cleanup()
		os.Exit(1)
	}
}

func runFail(cmdline string) {
	args := strings.Fields(cmdline)
	cmd := exec.Command(args[0], args[1:]...)
	out, err := cmd.CombinedOutput()
	if err == nil {
		fmt.Printf("$ %s\n", cmdline)
		fmt.Println(string(out))
		fmt.Println("SHOULD HAVE FAILED!")
		cleanup()
		os.Exit(1)
	}
}

func main() {
	// helloworld.go is package main
	run("go tool compile -o linkmain.o helloworld.go")
	run("go tool compile -pack -o linkmain.a helloworld.go")
	run("go tool link -o linkmain.exe linkmain.o")
	run("go tool link -o linkmain.exe linkmain.a")

	// linkmain.go is not
	run("go tool compile -o linkmain1.o linkmain.go")
	run("go tool compile -pack -o linkmain1.a linkmain.go")
	runFail("go tool link -o linkmain.exe linkmain1.o")
	runFail("go tool link -o linkmain.exe linkmain1.a")
	cleanup()
}
<|endoftext|>


<commit_before>// Package timelock implements a time based cache for dependencies
package timelock

import (
	"bytes"
	"encoding/json"
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/vube/depman/dep"
	"github.com/vube/depman/util"
)

const (
	// the filename of the cache json
	cacheFileName = ".depman.cache"

	// timeout in hours
	timeoutHours = 1.0
)

// flags
var (
	clear bool
	skip  bool
)

var (
	// the cache, a map of repos urls to timestamps
	cache map[string]time.Time

	// the full path to the cache json
	cacheFile string
)

func init() {
	flag.BoolVar(&clear, "clear-cache", false, "Delete the time based cache")
	flag.BoolVar(&skip, "skip-cache", false, "Skip the time based cache for this run only")
}

// Clear clears the cache, returns true if the cache was cleared
func Clear() (cleared bool) {
	cleared = clear

	if clear {
		parts := strings.Split(os.Getenv("GOPATH"), ":")
		cacheFile = filepath.Join(parts[0], cacheFileName)

		util.Print("Clearing cache file: " + cacheFile)

		_, err := os.Stat(cacheFile)
		if err != nil {
			return
		}

		err = os.Remove(cacheFile)
		if err != nil {
			util.Fatal(err)
		}
	}

	return
}

// Read reads the cache from disk
func Read() {
	parts := strings.Split(os.Getenv("GOPATH"), ":")
	cacheFile = filepath.Join(parts[0], cacheFileName)

	cache = make(map[string]time.Time)

	if !util.Exists(cacheFile) {
		return
	}

	util.Verbose("Reading cache file from " + cacheFile)

	data, err := ioutil.ReadFile(cacheFile)
	if err != nil {
		util.Fatal(err)
	}

	err = json.Unmarshal(data, &cache)
	if err != nil {
		return
	}
}

// Write writes the cache out to disk
func Write() {
	if skip {
		return
	}

	var buf bytes.Buffer
	str, err := json.Marshal(cache)
	if err == nil {
		json.Indent(&buf, str, "", "    ")
		data := []byte(buf.String() + "\n")

		util.Verbose("Writing cache file to " + cacheFile)

		err = ioutil.WriteFile(cacheFile, data, 0644)
		if err != nil {
			util.Fatal(err)
		}
	}
	return
}

// IsStale returns true if the cached dependency is older than timeoutHours
func IsStale(d *dep.Dependency) (stale bool) {

	if skip || d.SkipCache {
		return true
	}

	ts, ok := cache[d.Repo]

	// item is in the cache
	if ok {
		// item is old
		if time.Since(ts).Hours() > timeoutHours {
			stale = true
			cache[d.Repo] = time.Now()
		}
	} else {
		stale = true
		cache[d.Repo] = time.Now()
	}

	return
}
<commit_msg>timelock: Change clear cache message to yellow<commit_after>// Package timelock implements a time based cache for dependencies
package timelock

import (
	"bytes"
	"encoding/json"
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/vube/depman/colors"
	"github.com/vube/depman/dep"
	"github.com/vube/depman/util"
)

const (
	// the filename of the cache json
	cacheFileName = ".depman.cache"

	// timeout in hours
	timeoutHours = 1.0
)

// flags
var (
	clear bool
	skip  bool
)

var (
	// the cache, a map of repos urls to timestamps
	cache map[string]time.Time

	// the full path to the cache json
	cacheFile string
)

func init() {
	flag.BoolVar(&clear, "clear-cache", false, "Delete the time based cache")
	flag.BoolVar(&skip, "skip-cache", false, "Skip the time based cache for this run only")
}

// Clear clears the cache, returns true if the cache was cleared
func Clear() (cleared bool) {
	cleared = clear

	if clear {
		parts := strings.Split(os.Getenv("GOPATH"), ":")
		cacheFile = filepath.Join(parts[0], cacheFileName)

		util.Print(colors.Yellow("Clearing cache file: " + cacheFile))

		_, err := os.Stat(cacheFile)
		if err != nil {
			return
		}

		err = os.Remove(cacheFile)
		if err != nil {
			util.Fatal(err)
		}
	}

	return
}

// Read reads the cache from disk
func Read() {
	parts := strings.Split(os.Getenv("GOPATH"), ":")
	cacheFile = filepath.Join(parts[0], cacheFileName)

	cache = make(map[string]time.Time)

	if !util.Exists(cacheFile) {
		return
	}

	util.Verbose("Reading cache file from " + cacheFile)

	data, err := ioutil.ReadFile(cacheFile)
	if err != nil {
		util.Fatal(err)
	}

	err = json.Unmarshal(data, &cache)
	if err != nil {
		return
	}
}

// Write writes the cache out to disk
func Write() {
	if skip {
		return
	}

	var buf bytes.Buffer
	str, err := json.Marshal(cache)
	if err == nil {
		json.Indent(&buf, str, "", "    ")
		data := []byte(buf.String() + "\n")

		util.Verbose("Writing cache file to " + cacheFile)

		err = ioutil.WriteFile(cacheFile, data, 0644)
		if err != nil {
			util.Fatal(err)
		}
	}
	return
}

// IsStale returns true if the cached dependency is older than timeoutHours
func IsStale(d *dep.Dependency) (stale bool) {

	if skip || d.SkipCache {
		return true
	}

	ts, ok := cache[d.Repo]

	// item is in the cache
	if ok {
		// item is old
		if time.Since(ts).Hours() > timeoutHours {
			stale = true
			cache[d.Repo] = time.Now()
		}
	} else {
		stale = true
		cache[d.Repo] = time.Now()
	}

	return
}
<|endoftext|>


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

package cmd

import (
	"encoding/base64"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/TheThingsNetwork/ttn/core/types"
	"github.com/TheThingsNetwork/ttn/semtech"
	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/TheThingsNetwork/ttn/utils/pointer"
	"github.com/TheThingsNetwork/ttn/utils/random"
	"github.com/brocaar/lorawan"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// uplinkCmd represents the `uplink` command
var uplinkCmd = &cobra.Command{
	Use:   "uplink [ShouldConfirm] [DevAddr] [NwkSKey] [AppSKey] [Payload] [FCnt]",
	Short: "Send uplink messages to the network",
	Long:  `ttnctl uplink sends an uplink message to the network`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 6 {
			ctx.Fatalf("Insufficient arguments")
		}

		// Parse parameters
		var mtype lorawan.MType
		switch args[0] {
		case "yes":
			fallthrough
		case "true":
			mtype = lorawan.ConfirmedDataUp
		default:
			mtype = lorawan.UnconfirmedDataUp
		}

		devAddr, err := types.ParseDevAddr(args[1])
		if err != nil {
			ctx.Fatalf("Invalid DevAddr: %s", err)
		}

		nwkSKey, err := types.ParseNwkSKey(args[2])
		if err != nil {
			ctx.Fatalf("Invalid NwkSKey: %s", err)
		}

		appSKey, err := types.ParseAppSKey(args[3])
		if err != nil {
			ctx.Fatalf("Invalid appSKey: %s", err)
		}

		fcnt, err := strconv.ParseInt(args[5], 10, 64)
		if err != nil {
			ctx.Fatalf("Invalid FCnt: %s", err)
		}

		// Lorawan Payload
		macPayload := &lorawan.MACPayload{}
		macPayload.FHDR = lorawan.FHDR{
			DevAddr: lorawan.DevAddr(devAddr),
			FCnt:    uint32(fcnt),
		}
		macPayload.FPort = pointer.Uint8(1)
		if plain, _ := cmd.Flags().GetBool("plain"); plain {
			ctx.Warn("Sending data as plain text is bad practice. We recommend to transmit data in a binary format.")
			macPayload.FRMPayload = []lorawan.Payload{&lorawan.DataPayload{Bytes: []byte(args[4])}}
		} else {
			payload, err := util.ParseHEX(args[4], len(args[4]))
			if err != nil {
				ctx.Fatalf("Invalid hexadecimal payload. If you are trying to send a plain-text payload, use the --plain flag.")
			}
			macPayload.FRMPayload = []lorawan.Payload{&lorawan.DataPayload{Bytes: payload}}
		}
		phyPayload := &lorawan.PHYPayload{}
		phyPayload.MHDR = lorawan.MHDR{
			MType: mtype,
			Major: lorawan.LoRaWANR1,
		}
		phyPayload.MACPayload = macPayload
		if err := phyPayload.EncryptFRMPayload(lorawan.AES128Key(appSKey)); err != nil {
			ctx.Fatalf("Unable to encrypt frame payload: %s", err)
		}
		if err := phyPayload.SetMIC(lorawan.AES128Key(nwkSKey)); err != nil {
			ctx.Fatalf("Unable to set MIC: %s", err)
		}

		addr, err := net.ResolveUDPAddr("udp", viper.GetString("ttn-router"))
		if err != nil {
			ctx.Fatalf("Couldn't resolve UDP address: %s", err)
		}
		conn, err := net.DialUDP("udp", nil, addr)
		if err != nil {
			ctx.Fatalf("Couldn't Dial UDP connection: %s", err)
		}

		// Handle downlink
		chdown := make(chan bool)
		go func() {
			// Get PullAck
			buf := make([]byte, 1024)
			n, err := conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt := new(semtech.Packet)
			if err := pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received PullAck: %s", pkt)

			// Get Ack
			buf = make([]byte, 1024)
			n, err = conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt = new(semtech.Packet)
			if err := pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received Ack: %s", pkt)

			// Get Downlink, if any
			buf = make([]byte, 1024)
			n, err = conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt = new(semtech.Packet)
			if err = pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received Downlink: %s", pkt)
			defer func() { chdown <- true }()

			if pkt.Payload == nil || pkt.Payload.TXPK == nil || pkt.Payload.TXPK.Data == nil {
				ctx.Fatalf("No payload available in downlink response")
			}

			data, err := base64.RawStdEncoding.DecodeString(*pkt.Payload.TXPK.Data)
			if err != nil {
				ctx.Fatalf("Unable to decode data payload: %s", err)
			}

			payload := &lorawan.PHYPayload{}
			if err := payload.UnmarshalBinary(data); err != nil {
				ctx.Fatalf("Unable to retrieve LoRaWAN PhyPayload: %s", err)
			}

			micOK, _ := payload.ValidateMIC(lorawan.AES128Key(nwkSKey))
			if !micOK {
				ctx.Warn("MIC check failed.")
			}

			macPayload, ok := payload.MACPayload.(*lorawan.MACPayload)
			if !ok || len(macPayload.FRMPayload) > 1 {
				ctx.Fatalf("Unable to retrieve LoRaWAN MACPayload")
			}
			ctx.Infof("Frame counter: %d", macPayload.FHDR.FCnt)
			if len(macPayload.FRMPayload) > 0 {
				decrypted, err := lorawan.EncryptFRMPayload(
					lorawan.AES128Key(appSKey),
					false,
					lorawan.DevAddr(devAddr),
					macPayload.FHDR.FCnt,
					macPayload.FRMPayload[0].(*lorawan.DataPayload).Bytes,
				)
				if err != nil {
					ctx.Fatalf("Unable to decrypt MACPayload: %s", err)
				}
				if plain, _ := cmd.Flags().GetBool("plain"); plain {
					unprintable, _ := regexp.Compile(`[^[:print:]]`)
					if unprintable.Match(decrypted) {
						ctx.WithField("warning", "payload contains unprintable characters").Infof("Decrypted Payload: %X", decrypted)
					} else {
						ctx.Infof("%s", decrypted)
					}
				} else {
					ctx.Infof("Decrypted Payload: %X", decrypted)
				}
			} else {
				ctx.Infof("The frame payload was empty.")
			}
		}()

		// PULL_DATA Packet

		pullPacket := semtech.Packet{
			Version:    semtech.VERSION,
			GatewayId:  []byte{1, 2, 3, 4, 5, 6, 7, 8},
			Token:      []byte{1, 2},
			Identifier: semtech.PULL_DATA,
		}
		pullData, err := pullPacket.MarshalBinary()
		if err != nil {
			ctx.Fatal("Unable to construct pull_data")
		}

		// Router Packet
		data, err := phyPayload.MarshalBinary()
		if err != nil {
			ctx.Fatalf("Couldn't construct LoRaWAN physical payload: %s", err)
		}
		encoded := strings.Trim(base64.StdEncoding.EncodeToString(data), "=")
		payload := semtech.Packet{
			Identifier: semtech.PUSH_DATA,
			Token:      random.Token(),
			GatewayId:  []byte{1, 2, 3, 4, 5, 6, 7, 8},
			Version:    semtech.VERSION,
			Payload: &semtech.Payload{
				RXPK: []semtech.RXPK{
					{
						Rssi: pointer.Int32(random.Rssi()),
						Lsnr: pointer.Float32(random.Lsnr()),
						Freq: pointer.Float32(random.Freq()),
						Datr: pointer.String(random.Datr()),
						Codr: pointer.String(random.Codr()),
						Modu: pointer.String("LoRa"),
						Tmst: pointer.Uint32(1),
						Time: pointer.Time(time.Now().UTC()),
						Data: &encoded,
					},
				},
			},
		}

		ctx.Infof("Sending packet: %s", payload.String())

		data, err = payload.MarshalBinary()
		if err != nil {
			ctx.Fatalf("Unable to construct framepayload: %v", data)
		}

		_, err = conn.Write(pullData)
		if err != nil {
			ctx.Fatal("Unable to send pull_data")
		}

		_, err = conn.Write(data)
		if err != nil {
			ctx.Fatal("Unable to send payload")
		}

		select {
		case <-chdown:
		case <-time.After(2 * time.Second):
		}
	},
}

func init() {
	RootCmd.AddCommand(uplinkCmd)
	uplinkCmd.Flags().Bool("plain", false, "send payload as plain-text")
}
<commit_msg>Set stat to 1 in ttnctl uplink<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"encoding/base64"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/TheThingsNetwork/ttn/core/types"
	"github.com/TheThingsNetwork/ttn/semtech"
	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/TheThingsNetwork/ttn/utils/pointer"
	"github.com/TheThingsNetwork/ttn/utils/random"
	"github.com/brocaar/lorawan"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// uplinkCmd represents the `uplink` command
var uplinkCmd = &cobra.Command{
	Use:   "uplink [ShouldConfirm] [DevAddr] [NwkSKey] [AppSKey] [Payload] [FCnt]",
	Short: "Send uplink messages to the network",
	Long:  `ttnctl uplink sends an uplink message to the network`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 6 {
			ctx.Fatalf("Insufficient arguments")
		}

		// Parse parameters
		var mtype lorawan.MType
		switch args[0] {
		case "yes":
			fallthrough
		case "true":
			mtype = lorawan.ConfirmedDataUp
		default:
			mtype = lorawan.UnconfirmedDataUp
		}

		devAddr, err := types.ParseDevAddr(args[1])
		if err != nil {
			ctx.Fatalf("Invalid DevAddr: %s", err)
		}

		nwkSKey, err := types.ParseNwkSKey(args[2])
		if err != nil {
			ctx.Fatalf("Invalid NwkSKey: %s", err)
		}

		appSKey, err := types.ParseAppSKey(args[3])
		if err != nil {
			ctx.Fatalf("Invalid appSKey: %s", err)
		}

		fcnt, err := strconv.ParseInt(args[5], 10, 64)
		if err != nil {
			ctx.Fatalf("Invalid FCnt: %s", err)
		}

		// Lorawan Payload
		macPayload := &lorawan.MACPayload{}
		macPayload.FHDR = lorawan.FHDR{
			DevAddr: lorawan.DevAddr(devAddr),
			FCnt:    uint32(fcnt),
		}
		macPayload.FPort = pointer.Uint8(1)
		if plain, _ := cmd.Flags().GetBool("plain"); plain {
			ctx.Warn("Sending data as plain text is bad practice. We recommend to transmit data in a binary format.")
			macPayload.FRMPayload = []lorawan.Payload{&lorawan.DataPayload{Bytes: []byte(args[4])}}
		} else {
			payload, err := util.ParseHEX(args[4], len(args[4]))
			if err != nil {
				ctx.Fatalf("Invalid hexadecimal payload. If you are trying to send a plain-text payload, use the --plain flag.")
			}
			macPayload.FRMPayload = []lorawan.Payload{&lorawan.DataPayload{Bytes: payload}}
		}
		phyPayload := &lorawan.PHYPayload{}
		phyPayload.MHDR = lorawan.MHDR{
			MType: mtype,
			Major: lorawan.LoRaWANR1,
		}
		phyPayload.MACPayload = macPayload
		if err := phyPayload.EncryptFRMPayload(lorawan.AES128Key(appSKey)); err != nil {
			ctx.Fatalf("Unable to encrypt frame payload: %s", err)
		}
		if err := phyPayload.SetMIC(lorawan.AES128Key(nwkSKey)); err != nil {
			ctx.Fatalf("Unable to set MIC: %s", err)
		}

		addr, err := net.ResolveUDPAddr("udp", viper.GetString("ttn-router"))
		if err != nil {
			ctx.Fatalf("Couldn't resolve UDP address: %s", err)
		}
		conn, err := net.DialUDP("udp", nil, addr)
		if err != nil {
			ctx.Fatalf("Couldn't Dial UDP connection: %s", err)
		}

		// Handle downlink
		chdown := make(chan bool)
		go func() {
			// Get PullAck
			buf := make([]byte, 1024)
			n, err := conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt := new(semtech.Packet)
			if err := pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received PullAck: %s", pkt)

			// Get Ack
			buf = make([]byte, 1024)
			n, err = conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt = new(semtech.Packet)
			if err := pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received Ack: %s", pkt)

			// Get Downlink, if any
			buf = make([]byte, 1024)
			n, err = conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt = new(semtech.Packet)
			if err = pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received Downlink: %s", pkt)
			defer func() { chdown <- true }()

			if pkt.Payload == nil || pkt.Payload.TXPK == nil || pkt.Payload.TXPK.Data == nil {
				ctx.Fatalf("No payload available in downlink response")
			}

			data, err := base64.RawStdEncoding.DecodeString(*pkt.Payload.TXPK.Data)
			if err != nil {
				ctx.Fatalf("Unable to decode data payload: %s", err)
			}

			payload := &lorawan.PHYPayload{}
			if err := payload.UnmarshalBinary(data); err != nil {
				ctx.Fatalf("Unable to retrieve LoRaWAN PhyPayload: %s", err)
			}

			micOK, _ := payload.ValidateMIC(lorawan.AES128Key(nwkSKey))
			if !micOK {
				ctx.Warn("MIC check failed.")
			}

			macPayload, ok := payload.MACPayload.(*lorawan.MACPayload)
			if !ok || len(macPayload.FRMPayload) > 1 {
				ctx.Fatalf("Unable to retrieve LoRaWAN MACPayload")
			}
			ctx.Infof("Frame counter: %d", macPayload.FHDR.FCnt)
			if len(macPayload.FRMPayload) > 0 {
				decrypted, err := lorawan.EncryptFRMPayload(
					lorawan.AES128Key(appSKey),
					false,
					lorawan.DevAddr(devAddr),
					macPayload.FHDR.FCnt,
					macPayload.FRMPayload[0].(*lorawan.DataPayload).Bytes,
				)
				if err != nil {
					ctx.Fatalf("Unable to decrypt MACPayload: %s", err)
				}
				if plain, _ := cmd.Flags().GetBool("plain"); plain {
					unprintable, _ := regexp.Compile(`[^[:print:]]`)
					if unprintable.Match(decrypted) {
						ctx.WithField("warning", "payload contains unprintable characters").Infof("Decrypted Payload: %X", decrypted)
					} else {
						ctx.Infof("%s", decrypted)
					}
				} else {
					ctx.Infof("Decrypted Payload: %X", decrypted)
				}
			} else {
				ctx.Infof("The frame payload was empty.")
			}
		}()

		// PULL_DATA Packet

		pullPacket := semtech.Packet{
			Version:    semtech.VERSION,
			GatewayId:  []byte{1, 2, 3, 4, 5, 6, 7, 8},
			Token:      []byte{1, 2},
			Identifier: semtech.PULL_DATA,
		}
		pullData, err := pullPacket.MarshalBinary()
		if err != nil {
			ctx.Fatal("Unable to construct pull_data")
		}

		// Router Packet
		data, err := phyPayload.MarshalBinary()
		if err != nil {
			ctx.Fatalf("Couldn't construct LoRaWAN physical payload: %s", err)
		}
		encoded := strings.Trim(base64.StdEncoding.EncodeToString(data), "=")
		payload := semtech.Packet{
			Identifier: semtech.PUSH_DATA,
			Token:      random.Token(),
			GatewayId:  []byte{1, 2, 3, 4, 5, 6, 7, 8},
			Version:    semtech.VERSION,
			Payload: &semtech.Payload{
				RXPK: []semtech.RXPK{
					{
						Rssi: pointer.Int32(random.Rssi()),
						Lsnr: pointer.Float32(random.Lsnr()),
						Freq: pointer.Float32(random.Freq()),
						Datr: pointer.String(random.Datr()),
						Codr: pointer.String(random.Codr()),
						Modu: pointer.String("LoRa"),
						Tmst: pointer.Uint32(1),
						Time: pointer.Time(time.Now().UTC()),
						Data: &encoded,
						Stat: pointer.Int32(1),
					},
				},
			},
		}

		ctx.Infof("Sending packet: %s", payload.String())

		data, err = payload.MarshalBinary()
		if err != nil {
			ctx.Fatalf("Unable to construct framepayload: %v", data)
		}

		_, err = conn.Write(pullData)
		if err != nil {
			ctx.Fatal("Unable to send pull_data")
		}

		_, err = conn.Write(data)
		if err != nil {
			ctx.Fatal("Unable to send payload")
		}

		select {
		case <-chdown:
		case <-time.After(2 * time.Second):
		}
	},
}

func init() {
	RootCmd.AddCommand(uplinkCmd)
	uplinkCmd.Flags().Bool("plain", false, "send payload as plain-text")
}
<|endoftext|>


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

package txscript

import (
	"github.com/decred/dcrd/dcrec"
)

const (
	// MaxDataCarrierSize is the maximum number of bytes allowed in pushed
	// data to be considered a nulldata transaction.
	MaxDataCarrierSize = 256
)

// ScriptClass is an enumeration for the list of standard types of script.
type ScriptClass byte

// Classes of script payment known about in the blockchain.
const (
	NonStandardTy     ScriptClass = iota // None of the recognized forms.
	PubKeyTy                             // Pay pubkey.
	PubKeyHashTy                         // Pay pubkey hash.
	ScriptHashTy                         // Pay to script hash.
	MultiSigTy                           // Multi signature.
	NullDataTy                           // Empty data-only (provably prunable).
	StakeSubmissionTy                    // Stake submission.
	StakeGenTy                           // Stake generation
	StakeRevocationTy                    // Stake revocation.
	StakeSubChangeTy                     // Change for stake submission tx.
	PubkeyAltTy                          // Alternative signature pubkey.
	PubkeyHashAltTy                      // Alternative signature pubkey hash.
	TreasuryAddTy                        // Add value to treasury
	TreasuryGenTy                        // Generate utxos from treasury account
)

// scriptClassToName houses the human-readable strings which describe each
// script class.
var scriptClassToName = []string{
	NonStandardTy:     "nonstandard",
	PubKeyTy:          "pubkey",
	PubkeyAltTy:       "pubkeyalt",
	PubKeyHashTy:      "pubkeyhash",
	PubkeyHashAltTy:   "pubkeyhashalt",
	ScriptHashTy:      "scripthash",
	MultiSigTy:        "multisig",
	NullDataTy:        "nulldata",
	StakeSubmissionTy: "stakesubmission",
	StakeGenTy:        "stakegen",
	StakeRevocationTy: "stakerevoke",
	StakeSubChangeTy:  "sstxchange",
	TreasuryAddTy:     "treasuryadd",
	TreasuryGenTy:     "treasurygen",
}

// String implements the Stringer interface by returning the name of
// the enum script class. If the enum is invalid then "Invalid" will be
// returned.
func (t ScriptClass) String() string {
	if int(t) > len(scriptClassToName) || int(t) < 0 {
		return "Invalid"
	}
	return scriptClassToName[t]
}

// isStandardAltSignatureType returns whether or not the provided opcode
// represents a push of a standard alt signature type.
func isStandardAltSignatureType(op byte) bool {
	if !IsSmallInt(op) {
		return false
	}

	sigType := AsSmallInt(op)
	return sigType == dcrec.STEd25519 || sigType == dcrec.STSchnorrSecp256k1
}

// extractPubKeyHashAltDetails extracts the public key hash and signature type
// from the passed script if it is a standard pay-to-alt-pubkey-hash script.  It
// will return nil otherwise.
func extractPubKeyHashAltDetails(script []byte) ([]byte, dcrec.SignatureType) {
	// A pay-to-alt-pubkey-hash script is of the form:
	//  DUP HASH160 <20-byte hash> EQUALVERIFY SIGTYPE CHECKSIG
	//
	// The only two currently supported alternative signature types are ed25519
	// and schnorr + secp256k1 (with a compressed pubkey).
	//
	//  DUP HASH160 <20-byte hash> EQUALVERIFY <1-byte ed25519 sigtype> CHECKSIG
	//  DUP HASH160 <20-byte hash> EQUALVERIFY <1-byte schnorr+secp sigtype> CHECKSIG
	//
	//  Notice that OP_0 is not specified since signature type 0 disabled.

	if len(script) == 26 &&
		script[0] == OP_DUP &&
		script[1] == OP_HASH160 &&
		script[2] == OP_DATA_20 &&
		script[23] == OP_EQUALVERIFY &&
		isStandardAltSignatureType(script[24]) &&
		script[25] == OP_CHECKSIGALT {

		return script[3:23], dcrec.SignatureType(AsSmallInt(script[24]))
	}

	return nil, 0
}
<commit_msg>txscript: Remove unused extractPubKeyHashAltDetails.<commit_after>// Copyright (c) 2013-2017 The btcsuite developers
// Copyright (c) 2015-2021 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txscript

import (
	"github.com/decred/dcrd/dcrec"
)

const (
	// MaxDataCarrierSize is the maximum number of bytes allowed in pushed
	// data to be considered a nulldata transaction.
	MaxDataCarrierSize = 256
)

// ScriptClass is an enumeration for the list of standard types of script.
type ScriptClass byte

// Classes of script payment known about in the blockchain.
const (
	NonStandardTy     ScriptClass = iota // None of the recognized forms.
	PubKeyTy                             // Pay pubkey.
	PubKeyHashTy                         // Pay pubkey hash.
	ScriptHashTy                         // Pay to script hash.
	MultiSigTy                           // Multi signature.
	NullDataTy                           // Empty data-only (provably prunable).
	StakeSubmissionTy                    // Stake submission.
	StakeGenTy                           // Stake generation
	StakeRevocationTy                    // Stake revocation.
	StakeSubChangeTy                     // Change for stake submission tx.
	PubkeyAltTy                          // Alternative signature pubkey.
	PubkeyHashAltTy                      // Alternative signature pubkey hash.
	TreasuryAddTy                        // Add value to treasury
	TreasuryGenTy                        // Generate utxos from treasury account
)

// scriptClassToName houses the human-readable strings which describe each
// script class.
var scriptClassToName = []string{
	NonStandardTy:     "nonstandard",
	PubKeyTy:          "pubkey",
	PubkeyAltTy:       "pubkeyalt",
	PubKeyHashTy:      "pubkeyhash",
	PubkeyHashAltTy:   "pubkeyhashalt",
	ScriptHashTy:      "scripthash",
	MultiSigTy:        "multisig",
	NullDataTy:        "nulldata",
	StakeSubmissionTy: "stakesubmission",
	StakeGenTy:        "stakegen",
	StakeRevocationTy: "stakerevoke",
	StakeSubChangeTy:  "sstxchange",
	TreasuryAddTy:     "treasuryadd",
	TreasuryGenTy:     "treasurygen",
}

// String implements the Stringer interface by returning the name of
// the enum script class. If the enum is invalid then "Invalid" will be
// returned.
func (t ScriptClass) String() string {
	if int(t) > len(scriptClassToName) || int(t) < 0 {
		return "Invalid"
	}
	return scriptClassToName[t]
}

// isStandardAltSignatureType returns whether or not the provided opcode
// represents a push of a standard alt signature type.
func isStandardAltSignatureType(op byte) bool {
	if !IsSmallInt(op) {
		return false
	}

	sigType := AsSmallInt(op)
	return sigType == dcrec.STEd25519 || sigType == dcrec.STSchnorrSecp256k1
}
<|endoftext|>


<commit_before>package mount

// Type represents the type of a mount.
type Type string

const (
	// TypeBind BIND
	TypeBind Type = "bind"
	// TypeVolume VOLUME
	TypeVolume Type = "volume"
)

// Mount represents a mount (volume).
type Mount struct {
	Type     Type   `json:",omitempty"`
	Source   string `json:",omitempty"`
	Target   string `json:",omitempty"`
	ReadOnly bool   `json:",omitempty"`

	BindOptions   *BindOptions   `json:",omitempty"`
	VolumeOptions *VolumeOptions `json:",omitempty"`
}

// Propagation represents the propagation of a mount.
type Propagation string

const (
	// PropagationRPrivate RPRIVATE
	PropagationRPrivate Propagation = "rprivate"
	// PropagationPrivate PRIVATE
	PropagationPrivate Propagation = "private"
	// PropagationRShared RSHARED
	PropagationRShared Propagation = "rshared"
	// PropagationShared SHARED
	PropagationShared Propagation = "shared"
	// PropagationRSlave RSLAVE
	PropagationRSlave Propagation = "rslave"
	// PropagationSlave SLAVE
	PropagationSlave Propagation = "slave"
)

// BindOptions defines options specific to mounts of type "bind".
type BindOptions struct {
	Propagation Propagation `json:",omitempty"`
}

// VolumeOptions represents the options for a mount of type volume.
type VolumeOptions struct {
	Populate     bool              `json:",omitempty"`
	Labels       map[string]string `json:",omitempty"`
	DriverConfig *Driver           `json:",omitempty"`
}

// Driver represents a volume driver.
type Driver struct {
	Name    string            `json:",omitempty"`
	Options map[string]string `json:",omitempty"`
}
<commit_msg>Fix wrong field name on mount: s/Populate/NoCopy<commit_after>package mount

// Type represents the type of a mount.
type Type string

const (
	// TypeBind BIND
	TypeBind Type = "bind"
	// TypeVolume VOLUME
	TypeVolume Type = "volume"
)

// Mount represents a mount (volume).
type Mount struct {
	Type     Type   `json:",omitempty"`
	Source   string `json:",omitempty"`
	Target   string `json:",omitempty"`
	ReadOnly bool   `json:",omitempty"`

	BindOptions   *BindOptions   `json:",omitempty"`
	VolumeOptions *VolumeOptions `json:",omitempty"`
}

// Propagation represents the propagation of a mount.
type Propagation string

const (
	// PropagationRPrivate RPRIVATE
	PropagationRPrivate Propagation = "rprivate"
	// PropagationPrivate PRIVATE
	PropagationPrivate Propagation = "private"
	// PropagationRShared RSHARED
	PropagationRShared Propagation = "rshared"
	// PropagationShared SHARED
	PropagationShared Propagation = "shared"
	// PropagationRSlave RSLAVE
	PropagationRSlave Propagation = "rslave"
	// PropagationSlave SLAVE
	PropagationSlave Propagation = "slave"
)

// BindOptions defines options specific to mounts of type "bind".
type BindOptions struct {
	Propagation Propagation `json:",omitempty"`
}

// VolumeOptions represents the options for a mount of type volume.
type VolumeOptions struct {
	NoCopy       bool              `json:",omitempty"`
	Labels       map[string]string `json:",omitempty"`
	DriverConfig *Driver           `json:",omitempty"`
}

// Driver represents a volume driver.
type Driver struct {
	Name    string            `json:",omitempty"`
	Options map[string]string `json:",omitempty"`
}
<|endoftext|>


<commit_before>package xc

import (
	"bytes"
	"context"
	"crypto/rand"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/golang/protobuf/proto"
	"github.com/justwatchcom/gopass/backend/crypto/xc/keyring"
	"github.com/justwatchcom/gopass/backend/crypto/xc/xcpb"
	"github.com/stretchr/testify/assert"
)

type fakeAgent struct {
	pw string
}

func (f *fakeAgent) Ping() error {
	return nil
}

func (f *fakeAgent) Remove(string) error {
	return nil
}

func (f *fakeAgent) Passphrase(string, string) (string, error) {
	return f.pw, nil
}

func TestEncryptSimple(t *testing.T) {
	ctx := context.Background()

	td, err := ioutil.TempDir("", "gopass-")
	assert.NoError(t, err)
	defer func() {
		_ = os.RemoveAll(td)
	}()
	assert.NoError(t, os.Setenv("GOPASS_CONFIG", filepath.Join(td, ".gopass.yml")))
	assert.NoError(t, os.Setenv("GOPASS_HOMEDIR", td))

	passphrase := "test"

	k1, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)

	skr := keyring.NewSecring()
	assert.NoError(t, skr.Set(k1))

	pkr := keyring.NewPubring(skr)

	xc := &XC{
		pubring: pkr,
		secring: skr,
		client:  &fakeAgent{passphrase},
	}

	buf, err := xc.Encrypt(ctx, []byte("foobar"), []string{k1.Fingerprint()})
	assert.NoError(t, err)

	recps, err := xc.RecipientIDs(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, []string{k1.Fingerprint()}, recps)

	buf, err = xc.Decrypt(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, "foobar", string(buf))
}

func TestEncryptMultiKeys(t *testing.T) {
	ctx := context.Background()

	td, err := ioutil.TempDir("", "gopass-")
	assert.NoError(t, err)
	defer func() {
		_ = os.RemoveAll(td)
	}()
	assert.NoError(t, os.Setenv("GOPASS_CONFIG", filepath.Join(td, ".gopass.yml")))
	assert.NoError(t, os.Setenv("GOPASS_HOMEDIR", td))

	passphrase := "test"

	k1, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)
	k2, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)
	k3, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)

	skr := keyring.NewSecring()
	assert.NoError(t, skr.Set(k1))

	pkr := keyring.NewPubring(skr)
	assert.NoError(t, pkr.Set(&k2.PublicKey))
	assert.NoError(t, pkr.Set(&k3.PublicKey))

	xc := &XC{
		pubring: pkr,
		secring: skr,
		client:  &fakeAgent{passphrase},
	}

	buf, err := xc.Encrypt(ctx, []byte("foobar"), []string{k1.Fingerprint()})
	assert.NoError(t, err)

	recps, err := xc.RecipientIDs(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, []string{k1.Fingerprint()}, recps)

	buf, err = xc.Decrypt(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, "foobar", string(buf))
}

func TestEncryptChunks(t *testing.T) {
	ctx := context.Background()

	td, err := ioutil.TempDir("", "gopass-")
	assert.NoError(t, err)
	defer func() {
		_ = os.RemoveAll(td)
	}()
	assert.NoError(t, os.Setenv("GOPASS_CONFIG", filepath.Join(td, ".gopass.yml")))
	assert.NoError(t, os.Setenv("GOPASS_HOMEDIR", td))

	passphrase := "test"

	k1, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)

	skr := keyring.NewSecring()
	assert.NoError(t, skr.Set(k1))

	pkr := keyring.NewPubring(skr)

	xc := &XC{
		pubring: pkr,
		secring: skr,
		client:  &fakeAgent{passphrase},
	}

	plaintext := &bytes.Buffer{}
	p := make([]byte, 1024)
	for i := 0; i < 10*(chunkSizeMax/1024); i++ {
		_, _ = rand.Read(p)
		plaintext.Write(p)
	}
	assert.Equal(t, 10485760, plaintext.Len())

	buf, err := xc.Encrypt(ctx, plaintext.Bytes(), []string{k1.Fingerprint()})
	assert.NoError(t, err)

	// check recipients
	recps, err := xc.RecipientIDs(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, []string{k1.Fingerprint()}, recps)

	// check number of chunks
	msg := &xcpb.Message{}
	assert.NoError(t, proto.Unmarshal(buf, msg))
	assert.Equal(t, 10, len(msg.Chunks))

	// check decryption works and yields exactly the input
	buf, err = xc.Decrypt(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, plaintext.String(), string(buf))

}
<commit_msg>Add XC chunks reordering test case<commit_after>package xc

import (
	"bytes"
	"context"
	"crypto/rand"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/golang/protobuf/proto"
	"github.com/justwatchcom/gopass/backend/crypto/xc/keyring"
	"github.com/justwatchcom/gopass/backend/crypto/xc/xcpb"
	"github.com/stretchr/testify/assert"
)

type fakeAgent struct {
	pw string
}

func (f *fakeAgent) Ping() error {
	return nil
}

func (f *fakeAgent) Remove(string) error {
	return nil
}

func (f *fakeAgent) Passphrase(string, string) (string, error) {
	return f.pw, nil
}

func TestEncryptSimple(t *testing.T) {
	ctx := context.Background()

	td, err := ioutil.TempDir("", "gopass-")
	assert.NoError(t, err)
	defer func() {
		_ = os.RemoveAll(td)
	}()
	assert.NoError(t, os.Setenv("GOPASS_CONFIG", filepath.Join(td, ".gopass.yml")))
	assert.NoError(t, os.Setenv("GOPASS_HOMEDIR", td))

	passphrase := "test"

	k1, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)

	skr := keyring.NewSecring()
	assert.NoError(t, skr.Set(k1))

	pkr := keyring.NewPubring(skr)

	xc := &XC{
		pubring: pkr,
		secring: skr,
		client:  &fakeAgent{passphrase},
	}

	buf, err := xc.Encrypt(ctx, []byte("foobar"), []string{k1.Fingerprint()})
	assert.NoError(t, err)

	recps, err := xc.RecipientIDs(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, []string{k1.Fingerprint()}, recps)

	buf, err = xc.Decrypt(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, "foobar", string(buf))
}

func TestEncryptMultiKeys(t *testing.T) {
	ctx := context.Background()

	td, err := ioutil.TempDir("", "gopass-")
	assert.NoError(t, err)
	defer func() {
		_ = os.RemoveAll(td)
	}()
	assert.NoError(t, os.Setenv("GOPASS_CONFIG", filepath.Join(td, ".gopass.yml")))
	assert.NoError(t, os.Setenv("GOPASS_HOMEDIR", td))

	passphrase := "test"

	k1, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)
	k2, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)
	k3, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)

	skr := keyring.NewSecring()
	assert.NoError(t, skr.Set(k1))

	pkr := keyring.NewPubring(skr)
	assert.NoError(t, pkr.Set(&k2.PublicKey))
	assert.NoError(t, pkr.Set(&k3.PublicKey))

	xc := &XC{
		pubring: pkr,
		secring: skr,
		client:  &fakeAgent{passphrase},
	}

	buf, err := xc.Encrypt(ctx, []byte("foobar"), []string{k1.Fingerprint()})
	assert.NoError(t, err)

	recps, err := xc.RecipientIDs(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, []string{k1.Fingerprint()}, recps)

	buf, err = xc.Decrypt(ctx, buf)
	assert.NoError(t, err)
	assert.Equal(t, "foobar", string(buf))
}

func TestEncryptChunks(t *testing.T) {
	ctx := context.Background()

	td, err := ioutil.TempDir("", "gopass-")
	assert.NoError(t, err)
	defer func() {
		_ = os.RemoveAll(td)
	}()
	assert.NoError(t, os.Setenv("GOPASS_CONFIG", filepath.Join(td, ".gopass.yml")))
	assert.NoError(t, os.Setenv("GOPASS_HOMEDIR", td))

	passphrase := "test"

	k1, err := keyring.GenerateKeypair(passphrase)
	assert.NoError(t, err)

	skr := keyring.NewSecring()
	assert.NoError(t, skr.Set(k1))

	pkr := keyring.NewPubring(skr)

	xc := &XC{
		pubring: pkr,
		secring: skr,
		client:  &fakeAgent{passphrase},
	}

	plaintext := &bytes.Buffer{}
	p := make([]byte, 1024)
	for i := 0; i < 10*(chunkSizeMax/1024); i++ {
		_, _ = rand.Read(p)
		plaintext.Write(p)
	}
	assert.Equal(t, 10485760, plaintext.Len())

	ciphertext, err := xc.Encrypt(ctx, plaintext.Bytes(), []string{k1.Fingerprint()})
	assert.NoError(t, err)

	// check recipients
	recps, err := xc.RecipientIDs(ctx, ciphertext)
	assert.NoError(t, err)
	assert.Equal(t, []string{k1.Fingerprint()}, recps)

	// check number of chunks
	msg := &xcpb.Message{}
	assert.NoError(t, proto.Unmarshal(ciphertext, msg))
	assert.Equal(t, 10, len(msg.Chunks))

	// check decryption works and yields exactly the input
	plainagain, err := xc.Decrypt(ctx, ciphertext)
	assert.NoError(t, err)
	assert.Equal(t, plaintext.String(), string(plainagain))

	// reorder some chunks
	msg = &xcpb.Message{}
	assert.NoError(t, proto.Unmarshal(ciphertext, msg))
	assert.Equal(t, 10, len(msg.Chunks))

	msg.Chunks[0], msg.Chunks[1] = msg.Chunks[1], msg.Chunks[0]

	ciphertext, err = proto.Marshal(msg)
	assert.NoError(t, err)

	// check decryption fails
	_, err = xc.Decrypt(ctx, ciphertext)
	assert.Error(t, err)
}
<|endoftext|>


<commit_before>package routernode

import (
	"encoding/json"
	"fmt"
	"net"

	"github.com/Sirupsen/logrus"
	"github.com/jacksontj/dataman/src/router_node/metadata"
	"github.com/jacksontj/dataman/src/router_node/sharding"
	"github.com/jacksontj/dataman/src/storage_node"
	storagenodemetadata "github.com/jacksontj/dataman/src/storage_node/metadata"
)

func NewMetadataStore(config *storagenode.DatasourceInstanceConfig) (*MetadataStore, error) {
	// We want this layer to be responsible for initializing the storage node,
	// since this layer is responsible for the schema of the metadata anyways
	metaFunc, err := storagenodemetadata.StaticMetaFunc(schemaJson)
	if err != nil {
		return nil, err
	}

	store, err := config.GetStore(metaFunc)
	if err != nil {
		return nil, err
	}

	metaStore := &MetadataStore{
		Store: store,
	}

	return metaStore, nil
}

type MetadataStore struct {
	Store storagenode.StorageDataInterface
}

// TODO: split into get/list for each item?
// TODO: have error?
func (m *MetadataStore) GetMeta() *metadata.Meta {
	meta := metadata.NewMeta()

	// Add all nodes
	storageNodeResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "storage_node",
	})
	// TODO: better error handle
	if storageNodeResult.Error != "" {
		logrus.Fatalf("Error in getting storageNodeResult: %v", storageNodeResult.Error)
	}

	meta.Nodes = make(map[int64]*metadata.StorageNode)

	// for each database load the database + collections etc.
	for _, storageNodeRecord := range storageNodeResult.Return {
		meta.Nodes[storageNodeRecord["_id"].(int64)] = &metadata.StorageNode{
			ID:   storageNodeRecord["_id"].(int64),
			Name: storageNodeRecord["name"].(string),
			IP:   net.ParseIP(storageNodeRecord["ip"].(string)),
			Port: int(storageNodeRecord["port"].(int64)),
			// TODO: get the rest of it
			// Config
		}
	}

	// Load all of the datasource_instances
	datasourceInstanceResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "datasource_instance",
	})
	// TODO: better error handle
	if datasourceInstanceResult.Error != "" {
		logrus.Fatalf("Error in getting datasourceInstanceResult: %v", datasourceInstanceResult.Error)
	}
	for _, datasourceInstanceRecord := range datasourceInstanceResult.Return {
		datasourceInstance := metadata.NewDatasourceInstance(datasourceInstanceRecord["name"].(string))
		datasourceInstance.ID = datasourceInstanceRecord["_id"].(int64)
		datasourceInstance.StorageNodeID = datasourceInstanceRecord["storage_node_id"].(int64)
		datasourceInstance.StorageNode = meta.Nodes[datasourceInstanceRecord["storage_node_id"].(int64)]

		// Load all of the shard instances associated with this datasource_instance
		datasourceInstanceShardInstanceResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "datasource_instance_shard_instance",
			"filter": map[string]interface{}{
				"datasource_instance_id": datasourceInstanceRecord["_id"],
			},
		})
		// TODO: better error handle
		if datasourceInstanceShardInstanceResult.Error != "" {
			logrus.Fatalf("Error in getting datasourceInstanceShardInstanceResult: %v", datasourceInstanceShardInstanceResult.Error)
		}
		for _, datasourceInstanceShardInstanceRecord := range datasourceInstanceShardInstanceResult.Return {
			dsisi := &metadata.DatasourceInstanceShardInstance{
				ID: datasourceInstanceShardInstanceRecord["_id"].(int64),
				// TODO: need?
				//Name: datasourceInstanceShardInstanceRecord["name"].(string),
			}
			if databaseVShardID := datasourceInstanceShardInstanceRecord["database_vshard_instance_id"]; databaseVShardID != nil {
				datasourceInstance.DatabaseShards[dsisi.ID] = dsisi
			} else {
				datasourceInstance.CollectionShards[dsisi.ID] = dsisi
			}
		}

		// Set it in the map
		meta.DatasourceInstance[datasourceInstance.ID] = datasourceInstance
	}

	// Get all databases
	databaseResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "database",
	})
	// TODO: better error handle
	if databaseResult.Error != "" {
		logrus.Fatalf("Error in getting database list: %v", databaseResult.Error)
	}

	// for each database load the database + collections etc.
	for _, databaseRecord := range databaseResult.Return {
		database := metadata.NewDatabase(databaseRecord["name"].(string))
		database.ID = databaseRecord["_id"].(int64)

		// Load the database_vshards
		databaseVshardResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "database_vshard",
			"filter": map[string]interface{}{
				"database_id": databaseRecord["_id"],
			},
		})
		// TODO: better error handle
		if databaseVshardResult.Error != "" {
			logrus.Fatalf("Error in databaseVshardResult: %v", databaseVshardResult.Error)
		}

		databaseVshardRecord := databaseVshardResult.Return[0]
		database.VShard = metadata.NewDatabaseVShard()
		database.VShard.ID = databaseVshardRecord["_id"].(int64)
		database.VShard.ShardCount = databaseVshardRecord["shard_count"].(int64)
		database.Datastores = m.getDatastoreSetByDatabaseId(meta, databaseRecord["_id"].(int64))

		// TODO: order by!
		// Load all of the vshard instances
		databaseVshardInstanceResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "database_vshard_instance",
			"filter": map[string]interface{}{
				"database_vshard_id": databaseVshardRecord["_id"],
			},
		})
		// TODO: better error handle
		if databaseVshardInstanceResult.Error != "" {
			logrus.Fatalf("Error in databaseVshardInstanceResult: %v", databaseVshardInstanceResult.Error)
		}

		for _, databaseVshardInstanceRecord := range databaseVshardInstanceResult.Return {
			vshardInstance := &metadata.DatabaseVShardInstance{
				ID:             databaseVshardInstanceRecord["_id"].(int64),
				ShardInstance:  databaseVshardInstanceRecord["shard_instance"].(int64),
				DatastoreShard: meta.DatastoreShards[databaseVshardInstanceRecord["datastore_shard_id"].(int64)],
			}
			database.VShard.Instances = append(database.VShard.Instances, vshardInstance)
		}

		// TODO: resume here

		// Load all collections for the DB
		collectionResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "collection",
		})
		// TODO: better error handle
		if collectionResult.Error != "" {
			logrus.Fatalf("Error in collectionResult: %v", collectionResult.Error)
		}

		for _, collectionRecord := range collectionResult.Return {
			collection := metadata.NewCollection(collectionRecord["name"].(string))

			// TODO: load the rest of the collection

			// Load the partitions
			collectionPartitionResult := m.Store.Filter(map[string]interface{}{
				"db":             "dataman_router",
				"shard_instance": "public",
				"collection":     "collection_partition",
				"filter": map[string]interface{}{
					"collection_id": collectionRecord["_id"],
				},
			})
			// TODO: better error handle
			if collectionPartitionResult.Error != "" {
				logrus.Fatalf("Error in collectionPartitionResult: %v", collectionPartitionResult.Error)
			}

			collection.Partitions = make([]*metadata.CollectionPartition, len(collectionPartitionResult.Return))

			for i, collectionPartitionRecord := range collectionPartitionResult.Return {
				collection.Partitions[i] = &metadata.CollectionPartition{
					ID:      collectionPartitionRecord["_id"].(int64),
					StartId: collectionPartitionRecord["start_id"].(int64),
				}
				// EndId is optional (as this might be the first/only partition)
				if collectionPartitionRecord["end_id"] != nil {
					collection.Partitions[i].EndId = collectionPartitionRecord["end_id"].(int64)
				}

				// TODO: better
				shardConfigField := collectionPartitionRecord["shard_config_json"].(map[string]interface{})
				collection.Partitions[i].ShardConfig = &metadata.ShardConfig{
					Key:   shardConfigField["shard_key"].(string),
					Hash:  sharding.HashMethod(shardConfigField["hash_method"].(string)),
					Shard: sharding.ShardMethod(shardConfigField["shard_method"].(string)),
				}
				collection.Partitions[i].HashFunc = collection.Partitions[i].ShardConfig.Hash.Get()
				collection.Partitions[i].ShardFunc = collection.Partitions[i].ShardConfig.Shard.Get()
			}

			// Lastly add this collection to the database
			database.Collections[collection.Name] = collection
		}

		meta.Databases[database.Name] = database
	}

	return meta
}

// Here we want to query the database_datastore, and then get the datastores themselves
func (m *MetadataStore) getDatastoreSetByDatabaseId(meta *metadata.Meta, database_id int64) *metadata.DatastoreSet {
	set := metadata.NewDatastoreSet()

	// Get the datastore record
	databaseDatastoreResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "database_datastore",
		"filter": map[string]interface{}{
			"database_id": database_id,
		},
	})
	// TODO: better error handle
	if databaseDatastoreResult.Error != "" {
		logrus.Fatalf("Error in databaseDatastoreResult: %v", databaseDatastoreResult.Error)
	}

	for _, databaseDatastoreRecord := range databaseDatastoreResult.Return {
		datastore := m.getDatastoreById(meta, databaseDatastoreRecord["datastore_id"].(int64))

		// Set attributes associated with the linking table
		datastore.Read = databaseDatastoreRecord["read"].(bool)
		datastore.Write = databaseDatastoreRecord["write"].(bool)
		datastore.Required = databaseDatastoreRecord["required"].(bool)

		// Add to the set
		if datastore.Read {
			set.Read = append(set.Read, datastore)
		}

		if datastore.Write {
			if set.Write == nil {
				set.Write = datastore
			} else {
				logrus.Fatalf("Can only have one write datastore per database")
			}
		}

	}
	return set
}

// Get a single datastore by id
func (m *MetadataStore) getDatastoreById(meta *metadata.Meta, datastore_id int64) *metadata.Datastore {
	if datastore, ok := meta.Datastore[datastore_id]; ok {
		return datastore
	}
	datastoreResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "datastore",
		"filter": map[string]interface{}{
			"_id": datastore_id,
		},
	})
	// TODO: better error handle
	if datastoreResult.Error != "" {
		logrus.Fatalf("Error in datastoreResult: %v", datastoreResult.Error)
	}
	datastoreRecord := datastoreResult.Return[0]

	datastore := metadata.NewDatastore(datastoreRecord["name"].(string))
	datastore.ID = datastoreRecord["_id"].(int64)
	// TODO: define schema for shard config
	datastore.ShardConfig = datastoreRecord["shard_config_json"].(map[string]interface{})

	// TODO: order
	// Now load all the shards
	datastoreShardResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "datastore_shard",
		"filter": map[string]interface{}{
			"datastore_id": datastoreRecord["_id"],
		},
	})

	// TODO: better error handle
	if datastoreShardResult.Error != "" {
		logrus.Fatalf("Error in datastoreShardResult: %v", datastoreShardResult.Error)
	}

	for _, datastoreShardRecord := range datastoreShardResult.Return {
		datastoreShard := &metadata.DatastoreShard{
			ID:       datastoreShardRecord["_id"].(int64),
			Name:     datastoreShardRecord["name"].(string),
			Instance: datastoreShardRecord["shard_instance"].(int64),
			Replicas: metadata.NewDatastoreShardReplicaSet(),
		}

		// load all of the replicas
		datastoreShardReplicaResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "datastore_shard_replica",
			"filter": map[string]interface{}{
				"datastore_shard_id": datastoreShardRecord["_id"],
			},
		})

		// TODO: better error handle
		if datastoreShardReplicaResult.Error != "" {
			logrus.Fatalf("Error in datastoreShardReplicaResult: %v", datastoreShardReplicaResult.Error)
		}

		for _, datastoreShardReplicaRecord := range datastoreShardReplicaResult.Return {
			datastoreShardReplica := &metadata.DatastoreShardReplica{
				ID:         datastoreShardReplicaRecord["_id"].(int64),
				Master:     datastoreShardReplicaRecord["master"].(bool),
				Datasource: meta.DatasourceInstance[datastoreShardReplicaRecord["datasource_instance_id"].(int64)],
			}

			datastoreShard.Replicas.AddReplica(datastoreShardReplica)
		}
		datastore.Shards = append(datastore.Shards, datastoreShard)
		meta.DatastoreShards[datastoreShard.ID] = datastoreShard
	}

	meta.Datastore[datastore_id] = datastore
	return datastore
}

func structToRecord(item interface{}) map[string]interface{} {
	// TODO: better -- just don't want to spend all the time/space to do the conversions for now
	var record map[string]interface{}
	buf, _ := json.Marshal(item)
	json.Unmarshal(buf, &record)
	if _, ok := record["_id"]; ok {
		delete(record, "_id")
	}
	return record
}

func (m *MetadataStore) AddStorageNode(storageNode *metadata.StorageNode) error {
	record := structToRecord(storageNode)

	// load all of the replicas
	storageNodeResult := m.Store.Insert(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "storage_node",
		"record":         record,
	})

	// TODO: better error handle
	if storageNodeResult.Error != "" {
		return fmt.Errorf(storageNodeResult.Error)
	}

	return nil
}

func (m *MetadataStore) RemoveStorageNode(id int64) error {
	// load all of the replicas
	storageNodeResult := m.Store.Delete(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "storage_node",
		"_id":            id,
	})

	// TODO: better error handle
	if storageNodeResult.Error != "" {
		return fmt.Errorf(storageNodeResult.Error)
	}

	return nil
}
<commit_msg>Comments TODO<commit_after>package routernode

import (
	"encoding/json"
	"fmt"
	"net"

	"github.com/Sirupsen/logrus"
	"github.com/jacksontj/dataman/src/router_node/metadata"
	"github.com/jacksontj/dataman/src/router_node/sharding"
	"github.com/jacksontj/dataman/src/storage_node"
	storagenodemetadata "github.com/jacksontj/dataman/src/storage_node/metadata"
)

func NewMetadataStore(config *storagenode.DatasourceInstanceConfig) (*MetadataStore, error) {
	// We want this layer to be responsible for initializing the storage node,
	// since this layer is responsible for the schema of the metadata anyways
	metaFunc, err := storagenodemetadata.StaticMetaFunc(schemaJson)
	if err != nil {
		return nil, err
	}

	store, err := config.GetStore(metaFunc)
	if err != nil {
		return nil, err
	}

	metaStore := &MetadataStore{
		Store: store,
	}

	return metaStore, nil
}

type MetadataStore struct {
	Store storagenode.StorageDataInterface
}

// TODO: this should ideally load exactly *one* of any given record into a struct. This
// will require some work to do so, and we really should probably have something to codegen
// the record -> struct transition
// TODO: split into get/list for each item?
// TODO: have error?
func (m *MetadataStore) GetMeta() *metadata.Meta {
	meta := metadata.NewMeta()

	// Add all nodes
	storageNodeResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "storage_node",
	})
	// TODO: better error handle
	if storageNodeResult.Error != "" {
		logrus.Fatalf("Error in getting storageNodeResult: %v", storageNodeResult.Error)
	}

	meta.Nodes = make(map[int64]*metadata.StorageNode)

	// for each database load the database + collections etc.
	for _, storageNodeRecord := range storageNodeResult.Return {
		meta.Nodes[storageNodeRecord["_id"].(int64)] = &metadata.StorageNode{
			ID:   storageNodeRecord["_id"].(int64),
			Name: storageNodeRecord["name"].(string),
			IP:   net.ParseIP(storageNodeRecord["ip"].(string)),
			Port: int(storageNodeRecord["port"].(int64)),
			// TODO: get the rest of it
			// Config
		}
	}

	// Load all of the datasource_instances
	datasourceInstanceResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "datasource_instance",
	})
	// TODO: better error handle
	if datasourceInstanceResult.Error != "" {
		logrus.Fatalf("Error in getting datasourceInstanceResult: %v", datasourceInstanceResult.Error)
	}
	for _, datasourceInstanceRecord := range datasourceInstanceResult.Return {
		datasourceInstance := metadata.NewDatasourceInstance(datasourceInstanceRecord["name"].(string))
		datasourceInstance.ID = datasourceInstanceRecord["_id"].(int64)
		datasourceInstance.StorageNodeID = datasourceInstanceRecord["storage_node_id"].(int64)
		datasourceInstance.StorageNode = meta.Nodes[datasourceInstanceRecord["storage_node_id"].(int64)]

		// Load all of the shard instances associated with this datasource_instance
		datasourceInstanceShardInstanceResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "datasource_instance_shard_instance",
			"filter": map[string]interface{}{
				"datasource_instance_id": datasourceInstanceRecord["_id"],
			},
		})
		// TODO: better error handle
		if datasourceInstanceShardInstanceResult.Error != "" {
			logrus.Fatalf("Error in getting datasourceInstanceShardInstanceResult: %v", datasourceInstanceShardInstanceResult.Error)
		}
		for _, datasourceInstanceShardInstanceRecord := range datasourceInstanceShardInstanceResult.Return {
			dsisi := &metadata.DatasourceInstanceShardInstance{
				ID: datasourceInstanceShardInstanceRecord["_id"].(int64),
				// TODO: need?
				//Name: datasourceInstanceShardInstanceRecord["name"].(string),
			}
			if databaseVShardID := datasourceInstanceShardInstanceRecord["database_vshard_instance_id"]; databaseVShardID != nil {
				datasourceInstance.DatabaseShards[dsisi.ID] = dsisi
			} else {
				datasourceInstance.CollectionShards[dsisi.ID] = dsisi
			}
		}

		// Set it in the map
		meta.DatasourceInstance[datasourceInstance.ID] = datasourceInstance
	}

	// Get all databases
	databaseResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "database",
	})
	// TODO: better error handle
	if databaseResult.Error != "" {
		logrus.Fatalf("Error in getting database list: %v", databaseResult.Error)
	}

	// for each database load the database + collections etc.
	for _, databaseRecord := range databaseResult.Return {
		database := metadata.NewDatabase(databaseRecord["name"].(string))
		database.ID = databaseRecord["_id"].(int64)

		// Load the database_vshards
		databaseVshardResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "database_vshard",
			"filter": map[string]interface{}{
				"database_id": databaseRecord["_id"],
			},
		})
		// TODO: better error handle
		if databaseVshardResult.Error != "" {
			logrus.Fatalf("Error in databaseVshardResult: %v", databaseVshardResult.Error)
		}

		databaseVshardRecord := databaseVshardResult.Return[0]
		database.VShard = metadata.NewDatabaseVShard()
		database.VShard.ID = databaseVshardRecord["_id"].(int64)
		database.VShard.ShardCount = databaseVshardRecord["shard_count"].(int64)
		database.Datastores = m.getDatastoreSetByDatabaseId(meta, databaseRecord["_id"].(int64))

		// TODO: order by!
		// Load all of the vshard instances
		databaseVshardInstanceResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "database_vshard_instance",
			"filter": map[string]interface{}{
				"database_vshard_id": databaseVshardRecord["_id"],
			},
		})
		// TODO: better error handle
		if databaseVshardInstanceResult.Error != "" {
			logrus.Fatalf("Error in databaseVshardInstanceResult: %v", databaseVshardInstanceResult.Error)
		}

		for _, databaseVshardInstanceRecord := range databaseVshardInstanceResult.Return {
			vshardInstance := &metadata.DatabaseVShardInstance{
				ID:             databaseVshardInstanceRecord["_id"].(int64),
				ShardInstance:  databaseVshardInstanceRecord["shard_instance"].(int64),
				DatastoreShard: meta.DatastoreShards[databaseVshardInstanceRecord["datastore_shard_id"].(int64)],
			}
			database.VShard.Instances = append(database.VShard.Instances, vshardInstance)
		}

		// TODO: resume here

		// Load all collections for the DB
		collectionResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "collection",
		})
		// TODO: better error handle
		if collectionResult.Error != "" {
			logrus.Fatalf("Error in collectionResult: %v", collectionResult.Error)
		}

		for _, collectionRecord := range collectionResult.Return {
			collection := metadata.NewCollection(collectionRecord["name"].(string))

			// TODO: load the rest of the collection

			// Load the partitions
			collectionPartitionResult := m.Store.Filter(map[string]interface{}{
				"db":             "dataman_router",
				"shard_instance": "public",
				"collection":     "collection_partition",
				"filter": map[string]interface{}{
					"collection_id": collectionRecord["_id"],
				},
			})
			// TODO: better error handle
			if collectionPartitionResult.Error != "" {
				logrus.Fatalf("Error in collectionPartitionResult: %v", collectionPartitionResult.Error)
			}

			collection.Partitions = make([]*metadata.CollectionPartition, len(collectionPartitionResult.Return))

			for i, collectionPartitionRecord := range collectionPartitionResult.Return {
				collection.Partitions[i] = &metadata.CollectionPartition{
					ID:      collectionPartitionRecord["_id"].(int64),
					StartId: collectionPartitionRecord["start_id"].(int64),
				}
				// EndId is optional (as this might be the first/only partition)
				if collectionPartitionRecord["end_id"] != nil {
					collection.Partitions[i].EndId = collectionPartitionRecord["end_id"].(int64)
				}

				// TODO: better
				shardConfigField := collectionPartitionRecord["shard_config_json"].(map[string]interface{})
				collection.Partitions[i].ShardConfig = &metadata.ShardConfig{
					Key:   shardConfigField["shard_key"].(string),
					Hash:  sharding.HashMethod(shardConfigField["hash_method"].(string)),
					Shard: sharding.ShardMethod(shardConfigField["shard_method"].(string)),
				}
				collection.Partitions[i].HashFunc = collection.Partitions[i].ShardConfig.Hash.Get()
				collection.Partitions[i].ShardFunc = collection.Partitions[i].ShardConfig.Shard.Get()
			}

			// Lastly add this collection to the database
			database.Collections[collection.Name] = collection
		}

		meta.Databases[database.Name] = database
	}

	return meta
}

// Here we want to query the database_datastore, and then get the datastores themselves
func (m *MetadataStore) getDatastoreSetByDatabaseId(meta *metadata.Meta, database_id int64) *metadata.DatastoreSet {
	set := metadata.NewDatastoreSet()

	// Get the datastore record
	databaseDatastoreResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "database_datastore",
		"filter": map[string]interface{}{
			"database_id": database_id,
		},
	})
	// TODO: better error handle
	if databaseDatastoreResult.Error != "" {
		logrus.Fatalf("Error in databaseDatastoreResult: %v", databaseDatastoreResult.Error)
	}

	for _, databaseDatastoreRecord := range databaseDatastoreResult.Return {
		datastore := m.getDatastoreById(meta, databaseDatastoreRecord["datastore_id"].(int64))

		// Set attributes associated with the linking table
		datastore.Read = databaseDatastoreRecord["read"].(bool)
		datastore.Write = databaseDatastoreRecord["write"].(bool)
		datastore.Required = databaseDatastoreRecord["required"].(bool)

		// Add to the set
		if datastore.Read {
			set.Read = append(set.Read, datastore)
		}

		if datastore.Write {
			if set.Write == nil {
				set.Write = datastore
			} else {
				logrus.Fatalf("Can only have one write datastore per database")
			}
		}

	}
	return set
}

// Get a single datastore by id
func (m *MetadataStore) getDatastoreById(meta *metadata.Meta, datastore_id int64) *metadata.Datastore {
	if datastore, ok := meta.Datastore[datastore_id]; ok {
		return datastore
	}
	datastoreResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "datastore",
		"filter": map[string]interface{}{
			"_id": datastore_id,
		},
	})
	// TODO: better error handle
	if datastoreResult.Error != "" {
		logrus.Fatalf("Error in datastoreResult: %v", datastoreResult.Error)
	}
	datastoreRecord := datastoreResult.Return[0]

	datastore := metadata.NewDatastore(datastoreRecord["name"].(string))
	datastore.ID = datastoreRecord["_id"].(int64)
	// TODO: define schema for shard config
	datastore.ShardConfig = datastoreRecord["shard_config_json"].(map[string]interface{})

	// TODO: order
	// Now load all the shards
	datastoreShardResult := m.Store.Filter(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "datastore_shard",
		"filter": map[string]interface{}{
			"datastore_id": datastoreRecord["_id"],
		},
	})

	// TODO: better error handle
	if datastoreShardResult.Error != "" {
		logrus.Fatalf("Error in datastoreShardResult: %v", datastoreShardResult.Error)
	}

	for _, datastoreShardRecord := range datastoreShardResult.Return {
		datastoreShard := &metadata.DatastoreShard{
			ID:       datastoreShardRecord["_id"].(int64),
			Name:     datastoreShardRecord["name"].(string),
			Instance: datastoreShardRecord["shard_instance"].(int64),
			Replicas: metadata.NewDatastoreShardReplicaSet(),
		}

		// load all of the replicas
		datastoreShardReplicaResult := m.Store.Filter(map[string]interface{}{
			"db":             "dataman_router",
			"shard_instance": "public",
			"collection":     "datastore_shard_replica",
			"filter": map[string]interface{}{
				"datastore_shard_id": datastoreShardRecord["_id"],
			},
		})

		// TODO: better error handle
		if datastoreShardReplicaResult.Error != "" {
			logrus.Fatalf("Error in datastoreShardReplicaResult: %v", datastoreShardReplicaResult.Error)
		}

		for _, datastoreShardReplicaRecord := range datastoreShardReplicaResult.Return {
			datastoreShardReplica := &metadata.DatastoreShardReplica{
				ID:         datastoreShardReplicaRecord["_id"].(int64),
				Master:     datastoreShardReplicaRecord["master"].(bool),
				Datasource: meta.DatasourceInstance[datastoreShardReplicaRecord["datasource_instance_id"].(int64)],
			}

			datastoreShard.Replicas.AddReplica(datastoreShardReplica)
		}
		datastore.Shards = append(datastore.Shards, datastoreShard)
		meta.DatastoreShards[datastoreShard.ID] = datastoreShard
	}

	meta.Datastore[datastore_id] = datastore
	return datastore
}

func structToRecord(item interface{}) map[string]interface{} {
	// TODO: better -- just don't want to spend all the time/space to do the conversions for now
	var record map[string]interface{}
	buf, _ := json.Marshal(item)
	json.Unmarshal(buf, &record)
	if _, ok := record["_id"]; ok {
		delete(record, "_id")
	}
	return record
}

func (m *MetadataStore) AddStorageNode(storageNode *metadata.StorageNode) error {
	record := structToRecord(storageNode)

	// load all of the replicas
	storageNodeResult := m.Store.Insert(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "storage_node",
		"record":         record,
	})

	// TODO: better error handle
	if storageNodeResult.Error != "" {
		return fmt.Errorf(storageNodeResult.Error)
	}

	return nil
}

func (m *MetadataStore) RemoveStorageNode(id int64) error {
	// load all of the replicas
	storageNodeResult := m.Store.Delete(map[string]interface{}{
		"db":             "dataman_router",
		"shard_instance": "public",
		"collection":     "storage_node",
		"_id":            id,
	})

	// TODO: better error handle
	if storageNodeResult.Error != "" {
		return fmt.Errorf(storageNodeResult.Error)
	}

	return nil
}
<|endoftext|>