<commit_before>package api

import (
	"context"
	"net/http"

	"github.com/gorilla/mux"

	"github.com/ovh/cds/engine/api/project"
	"github.com/ovh/cds/engine/api/workflow"
	"github.com/ovh/cds/sdk"
)

func (api *API) getWorkflowTriggerConditionHandler() Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)
		key := vars["permProjectKey"]
		name := vars["workflowName"]

		id, errID := requestVarInt(r, "nodeID")
		if errID != nil {
			return errID
		}

		proj, errproj := project.Load(api.mustDB(), api.Cache, key, getUser(ctx), project.LoadOptions.WithVariables)
		if errproj != nil {
			return sdk.WrapError(errproj, "getWorkflowTriggerConditionHandler> Unable to load project")
		}

		wf, errw := workflow.Load(api.mustDB(), api.Cache, key, name, getUser(ctx))
		if errw != nil {
			return sdk.WrapError(errw, "getWorkflowTriggerConditionHandler> Unable to load workflow")
		}

		wr, errr := workflow.LoadLastRun(api.mustDB(), key, name)
		if errr != nil {
			if errr != sdk.ErrWorkflowNotFound {
				return sdk.WrapError(errr, "getWorkflowTriggerConditionHandler> Unable to load last run workflow")
			}
		}

		params, errp := workflow.NodeBuildParameters(proj, wf, wr, id, getUser(ctx))
		if errp != nil {
			return sdk.WrapError(errr, "getWorkflowTriggerConditionHandler> Unable to load build parameters")
		}

		data := struct {
			Operators      map[string]string `json:"operators"`
			ConditionNames []string          `json:"names"`
		}{
			Operators: sdk.WorkflowConditionsOperators,
		}

		var statusParamFound bool
		for _, p := range params {
			if p.Name == "cds.status" {
				statusParamFound = true
			}
			data.ConditionNames = append(data.ConditionNames, p.Name)
		}

		if !statusParamFound {
			data.ConditionNames = append(data.ConditionNames, "cds.status")
		}

		data.ConditionNames = append(data.ConditionNames, "cds.dest.pipeline")

		refNode := wf.GetNode(id)
		if refNode == nil {
			return sdk.WrapError(sdk.ErrWorkflowNodeNotFound, "getWorkflowTriggerConditionHandler> Unable to load workflow node")
		}
		if refNode.Context != nil && refNode.Context.Application != nil {
			data.ConditionNames = append(data.ConditionNames, "cds.dest.application")
		}
		if refNode.Context != nil && refNode.Context.Environment != nil {
			data.ConditionNames = append(data.ConditionNames, "cds.dest.environment")
		}

		ancestorIds := refNode.Ancestors(wf)
		for _, aID := range ancestorIds {
			ancestor := wf.GetNode(aID)
			if ancestor == nil {
				continue
			}
			var found bool
			for _, s := range data.ConditionNames {
				if s == "workflow."+ancestor.Name+".status" {
					found = true
				}
			}
			if !found {
				data.ConditionNames = append(data.ConditionNames, "workflow."+ancestor.Name+".status")
			}
		}

		return WriteJSON(w, r, data, http.StatusOK)
	}
}

func (api *API) getWorkflowTriggerJoinConditionHandler() Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)
		key := vars["permProjectKey"]
		name := vars["workflowName"]

		id, errID := requestVarInt(r, "joinID")
		if errID != nil {
			return errID
		}

		proj, errproj := project.Load(api.mustDB(), api.Cache, key, getUser(ctx), project.LoadOptions.WithVariables)
		if errproj != nil {
			return sdk.WrapError(errproj, "getWorkflowTriggerConditionHandler> Unable to load project")
		}

		wf, errw := workflow.Load(api.mustDB(), api.Cache, key, name, getUser(ctx))
		if errw != nil {
			return sdk.WrapError(errw, "getWorkflowTriggerJoinConditionHandler> Unable to load workflow")
		}

		wr, errr := workflow.LoadLastRun(api.mustDB(), key, name)
		if errr != nil {
			if errr != sdk.ErrWorkflowNotFound {
				return sdk.WrapError(errr, "getWorkflowTriggerJoinConditionHandler> Unable to load last run workflow")
			}
		}

		j := wf.GetJoin(id)
		if j == nil {
			return sdk.ErrWorkflowNodeJoinNotFound
		}

		data := struct {
			Operators      map[string]string `json:"operators"`
			ConditionNames []string          `json:"names"`
		}{
			Operators: sdk.WorkflowConditionsOperators,
		}

		allparams := map[string]string{}
		for _, i := range j.SourceNodeIDs {
			params, errp := workflow.NodeBuildParameters(proj, wf, wr, i, getUser(ctx))
			if errp != nil {
				return sdk.WrapError(errr, "getWorkflowTriggerJoinConditionHandler> Unable to load build parameters")
			}
			allparams = sdk.ParametersMapMerge(allparams, sdk.ParametersToMap(params))
		}

		for k := range allparams {
			data.ConditionNames = append(data.ConditionNames, k)
		}

		return WriteJSON(w, r, data, http.StatusOK)
	}
}
<commit_msg>feat (api): Improve workflow join trigger condition (#1190)<commit_after>package api

import (
	"context"
	"net/http"

	"github.com/gorilla/mux"

	"github.com/ovh/cds/engine/api/project"
	"github.com/ovh/cds/engine/api/workflow"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

func (api *API) getWorkflowTriggerConditionHandler() Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)
		key := vars["permProjectKey"]
		name := vars["workflowName"]

		id, errID := requestVarInt(r, "nodeID")
		if errID != nil {
			return errID
		}

		proj, errproj := project.Load(api.mustDB(), api.Cache, key, getUser(ctx), project.LoadOptions.WithVariables)
		if errproj != nil {
			return sdk.WrapError(errproj, "getWorkflowTriggerConditionHandler> Unable to load project")
		}

		wf, errw := workflow.Load(api.mustDB(), api.Cache, key, name, getUser(ctx))
		if errw != nil {
			return sdk.WrapError(errw, "getWorkflowTriggerConditionHandler> Unable to load workflow")
		}

		wr, errr := workflow.LoadLastRun(api.mustDB(), key, name)
		if errr != nil {
			if errr != sdk.ErrWorkflowNotFound {
				return sdk.WrapError(errr, "getWorkflowTriggerConditionHandler> Unable to load last run workflow")
			}
		}

		params, errp := workflow.NodeBuildParameters(proj, wf, wr, id, getUser(ctx))
		if errp != nil {
			return sdk.WrapError(errr, "getWorkflowTriggerConditionHandler> Unable to load build parameters")
		}

		data := struct {
			Operators      map[string]string `json:"operators"`
			ConditionNames []string          `json:"names"`
		}{
			Operators: sdk.WorkflowConditionsOperators,
		}

		var statusParamFound bool
		for _, p := range params {
			if p.Name == "cds.status" {
				statusParamFound = true
			}
			data.ConditionNames = append(data.ConditionNames, p.Name)
		}

		if !statusParamFound {
			data.ConditionNames = append(data.ConditionNames, "cds.status")
		}

		data.ConditionNames = append(data.ConditionNames, "cds.dest.pipeline")

		refNode := wf.GetNode(id)
		if refNode == nil {
			return sdk.WrapError(sdk.ErrWorkflowNodeNotFound, "getWorkflowTriggerConditionHandler> Unable to load workflow node")
		}
		if refNode.Context != nil && refNode.Context.Application != nil {
			data.ConditionNames = append(data.ConditionNames, "cds.dest.application")
		}
		if refNode.Context != nil && refNode.Context.Environment != nil {
			data.ConditionNames = append(data.ConditionNames, "cds.dest.environment")
		}

		ancestorIds := refNode.Ancestors(wf)
		for _, aID := range ancestorIds {
			ancestor := wf.GetNode(aID)
			if ancestor == nil {
				continue
			}
			var found bool
			for _, s := range data.ConditionNames {
				if s == "workflow."+ancestor.Name+".status" {
					found = true
				}
			}
			if !found {
				data.ConditionNames = append(data.ConditionNames, "workflow."+ancestor.Name+".status")
			}
		}

		return WriteJSON(w, r, data, http.StatusOK)
	}
}

func (api *API) getWorkflowTriggerJoinConditionHandler() Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)
		key := vars["permProjectKey"]
		name := vars["workflowName"]

		id, errID := requestVarInt(r, "joinID")
		if errID != nil {
			return errID
		}

		proj, errproj := project.Load(api.mustDB(), api.Cache, key, getUser(ctx), project.LoadOptions.WithVariables)
		if errproj != nil {
			return sdk.WrapError(errproj, "getWorkflowTriggerConditionHandler> Unable to load project")
		}

		wf, errw := workflow.Load(api.mustDB(), api.Cache, key, name, getUser(ctx))
		if errw != nil {
			return sdk.WrapError(errw, "getWorkflowTriggerJoinConditionHandler> Unable to load workflow")
		}

		wr, errr := workflow.LoadLastRun(api.mustDB(), key, name)
		if errr != nil {
			if errr != sdk.ErrWorkflowNotFound {
				return sdk.WrapError(errr, "getWorkflowTriggerJoinConditionHandler> Unable to load last run workflow")
			}
		}

		j := wf.GetJoin(id)
		if j == nil {
			return sdk.ErrWorkflowNodeJoinNotFound
		}

		data := struct {
			Operators      map[string]string `json:"operators"`
			ConditionNames []string          `json:"names"`
		}{
			Operators: sdk.WorkflowConditionsOperators,
		}

		//First we merge all build parameters from all source nodes
		allparams := map[string]string{}
		for _, i := range j.SourceNodeIDs {
			params, errp := workflow.NodeBuildParameters(proj, wf, wr, i, getUser(ctx))
			if errp != nil {
				return sdk.WrapError(errr, "getWorkflowTriggerJoinConditionHandler> Unable to load build parameters")
			}
			allparams = sdk.ParametersMapMerge(allparams, sdk.ParametersToMap(params))
		}
		for k := range allparams {
			data.ConditionNames = append(data.ConditionNames, k)
		}

		//Then we push all ancestors status if needed
		for _, i := range j.SourceNodeIDs {
			refNode := wf.GetNode(i)
			if refNode == nil {
				log.Error("getWorkflowTriggerJoinConditionHandler> Unable to get node %d", i)
				continue
			}
			var found bool
			for _, s := range data.ConditionNames {
				if s == "workflow."+refNode.Name+".status" {
					found = true
				}
			}
			if !found {
				data.ConditionNames = append(data.ConditionNames, "workflow."+refNode.Name+".status")
			}
			ancestorIds := refNode.Ancestors(wf)
			for _, aID := range ancestorIds {
				ancestor := wf.GetNode(aID)
				if ancestor == nil {
					continue
				}
				var found bool
				for _, s := range data.ConditionNames {
					if s == "workflow."+ancestor.Name+".status" {
						found = true
					}
				}
				if !found {
					data.ConditionNames = append(data.ConditionNames, "workflow."+ancestor.Name+".status")
				}
			}
		}

		return WriteJSON(w, r, data, http.StatusOK)
	}
}
<|endoftext|>


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

package azure

import (
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	. "launchpad.net/gocheck"
	"launchpad.net/gwacl"
	"launchpad.net/juju-core/errors"
)

type storageSuite struct {
	providerSuite
}

var _ = Suite(&storageSuite{})

func makeResponse(content string, status int) *http.Response {
	return &http.Response{
		Status:     fmt.Sprintf("%d", status),
		StatusCode: status,
		Body:       ioutil.NopCloser(strings.NewReader(content)),
	}
}

// MockingTransportExchange is a recording of a request and a response over
// HTTP.
type MockingTransportExchange struct {
	Request  *http.Request
	Response *http.Response
	Error    error
}

// MockingTransport is used as an http.Client.Transport for testing.  It
// records the sequence of requests, and returns a predetermined sequence of
// Responses and errors.
type MockingTransport struct {
	Exchanges     []*MockingTransportExchange
	ExchangeCount int
}

// MockingTransport implements the http.RoundTripper interface.
var _ http.RoundTripper = &MockingTransport{}

func (t *MockingTransport) AddExchange(response *http.Response, err error) {
	exchange := MockingTransportExchange{Response: response, Error: err}
	t.Exchanges = append(t.Exchanges, &exchange)
}

func (t *MockingTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	exchange := t.Exchanges[t.ExchangeCount]
	t.ExchangeCount++
	exchange.Request = req
	return exchange.Response, exchange.Error
}

// testStorageContext is a struct implementing the storageContext interface
// used in test.  It will return, via getContainer() and getStorageContext()
// the objects used at creation time.
type testStorageContext struct {
	container      string
	storageContext *gwacl.StorageContext
}

func (context *testStorageContext) getContainer() string {
	return context.container
}

func (context *testStorageContext) getStorageContext() (*gwacl.StorageContext, error) {
	return context.storageContext, nil
}

// makeFakeStorage creates a test azureStorage object that will talk to a
// fake HTTP server set up to always return preconfigured http.Response objects.
// The MockingTransport object can be used to check that the expected query has
// been issued to the test server.
func makeFakeStorage(container, account string) (azureStorage, *MockingTransport) {
	transport := &MockingTransport{}
	client := &http.Client{Transport: transport}
	storageContext := gwacl.NewTestStorageContext(client)
	storageContext.Account = account
	context := &testStorageContext{container: container, storageContext: storageContext}
	azStorage := azureStorage{storageContext: context}
	return azStorage, transport
}

var blobListResponse = `
  <?xml version="1.0" encoding="utf-8"?>
  <EnumerationResults ContainerName="http://myaccount.blob.core.windows.net/mycontainer">
    <Prefix>prefix</Prefix>
    <Marker>marker</Marker>
    <MaxResults>maxresults</MaxResults>
    <Delimiter>delimiter</Delimiter>
    <Blobs>
      <Blob>
        <Name>prefix-1</Name>
        <Url>blob-url1</Url>
      </Blob>
      <Blob>
        <Name>prefix-2</Name>
        <Url>blob-url2</Url>
      </Blob>
    </Blobs>
    <NextMarker />
  </EnumerationResults>`

func (*storageSuite) TestList(c *C) {
	container := "container"
	response := makeResponse(blobListResponse, http.StatusOK)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)

	prefix := "prefix"
	names, err := azStorage.List(prefix)
	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	// The prefix has been passed down as a query parameter.
	c.Check(transport.Exchanges[0].Request.URL.Query()["prefix"], DeepEquals, []string{prefix})
	// The container name is used in the requested URL.
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, ".*"+container+".*")
	c.Check(names, DeepEquals, []string{"prefix-1", "prefix-2"})
}

func (*storageSuite) TestListWithNonexistentContainerReturnsNoFiles(c *C) {
	// If Azure returns a 404 it means the container doesn't exist. In this
	// case the provider should interpret this as "no files" and return nil.
	container := "container"
	response := makeResponse("", http.StatusNotFound)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)

	names, err := azStorage.List("prefix")
	c.Assert(err, IsNil)
	c.Assert(names, IsNil)
}

func (*storageSuite) TestGet(c *C) {
	blobContent := "test blob"
	container := "container"
	filename := "blobname"
	response := makeResponse(blobContent, http.StatusOK)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)

	reader, err := azStorage.Get(filename)
	c.Assert(err, IsNil)
	c.Assert(reader, NotNil)
	defer reader.Close()

	context, err := azStorage.getStorageContext()
	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, context.GetFileURL(container, filename)+"?.*")
	data, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(string(data), Equals, blobContent)
}

func (*storageSuite) TestGetReturnsNotFoundIf404(c *C) {
	container := "container"
	filename := "blobname"
	response := makeResponse("not found", http.StatusNotFound)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)
	_, err := azStorage.Get(filename)
	c.Assert(err, NotNil)
	c.Check(errors.IsNotFoundError(err), Equals, true)
}

func (*storageSuite) TestPut(c *C) {
	blobContent := "test blob"
	container := "container"
	filename := "blobname"
	azStorage, transport := makeFakeStorage(container, "account")
	// The create container call makes two exchanges.
	transport.AddExchange(makeResponse("", http.StatusNotFound), nil)
	transport.AddExchange(makeResponse("", http.StatusCreated), nil)
	putResponse := makeResponse("", http.StatusCreated)
	transport.AddExchange(putResponse, nil)
	transport.AddExchange(putResponse, nil)
	err := azStorage.Put(filename, strings.NewReader(blobContent), int64(len(blobContent)))
	c.Assert(err, IsNil)

	context, err := azStorage.getStorageContext()
	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 4)
	c.Check(transport.Exchanges[2].Request.URL.String(), Matches, context.GetFileURL(container, filename)+"?.*")
}

func (*storageSuite) TestRemove(c *C) {
	container := "container"
	filename := "blobname"
	response := makeResponse("", http.StatusAccepted)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)
	err := azStorage.Remove(filename)
	c.Assert(err, IsNil)

	context, err := azStorage.getStorageContext()
	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, context.GetFileURL(container, filename)+"?.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "DELETE")
}

func (*storageSuite) TestRemoveErrors(c *C) {
	container := "container"
	filename := "blobname"
	response := makeResponse("", http.StatusForbidden)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)
	err := azStorage.Remove(filename)
	c.Assert(err, NotNil)
}

func (*storageSuite) TestRemoveAll(c *C) {
	// When we ask gwacl to remove all blobs, it calls DeleteContainer.
	response := makeResponse("", http.StatusAccepted)
	storage, transport := makeFakeStorage("cntnr", "account")
	transport.AddExchange(response, nil)

	err := storage.RemoveAll()
	c.Assert(err, IsNil)

	_, err = storage.getStorageContext()
	c.Assert(err, IsNil)
	// Without going too far into gwacl's innards, this is roughly what
	// it needs to do in order to delete a container.
	c.Assert(transport.ExchangeCount, Equals, 1)
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "DELETE")
}

func (*storageSuite) TestRemoveNonExistentBlobSucceeds(c *C) {
	container := "container"
	filename := "blobname"
	response := makeResponse("", http.StatusNotFound)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)
	err := azStorage.Remove(filename)
	c.Assert(err, IsNil)
}

func (*storageSuite) TestURL(c *C) {
	container := "container"
	filename := "blobname"
	account := "account"
	azStorage, _ := makeFakeStorage(container, account)
	URL, err := azStorage.URL(filename)
	c.Assert(err, IsNil)
	parsedURL, err := url.Parse(URL)
	c.Assert(err, IsNil)
	c.Check(parsedURL.Host, Matches, fmt.Sprintf("%s.blob.core.windows.net", account))
	c.Check(parsedURL.Path, Matches, fmt.Sprintf("/%s/%s", container, filename))
	values, err := url.ParseQuery(parsedURL.RawQuery)
	c.Assert(err, IsNil)
	signature := values.Get("sig")
	// The query string contains a non-empty signature.
	c.Check(signature, Not(HasLen), 0)
	// The signature is base64-encoded.
	_, err = base64.StdEncoding.DecodeString(signature)
	c.Assert(err, IsNil)
}

func (*storageSuite) TestCreateContainerCreatesContainerIfDoesNotExist(c *C) {
	azStorage, transport := makeFakeStorage("", "account")
	transport.AddExchange(makeResponse("", http.StatusNotFound), nil)
	transport.AddExchange(makeResponse("", http.StatusCreated), nil)

	err := azStorage.createContainer("cntnr")

	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 2)
	// Without going too far into gwacl's innards, this is roughly what
	// it needs to do in order to call GetContainerProperties.
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "GET")

	// ... and for CreateContainer.
	c.Check(transport.Exchanges[1].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[1].Request.Method, Equals, "PUT")
}

func (*storageSuite) TestCreateContainerIsDoneIfContainerAlreadyExists(c *C) {
	container := ""
	azStorage, transport := makeFakeStorage(container, "account")
	header := make(http.Header)
	header.Add("Last-Modified", "last-modified")
	header.Add("ETag", "etag")
	header.Add("X-Ms-Lease-Status", "status")
	header.Add("X-Ms-Lease-State", "state")
	header.Add("X-Ms-Lease-Duration", "duration")
	response := makeResponse("", http.StatusOK)
	response.Header = header
	transport.AddExchange(response, nil)

	err := azStorage.createContainer("cntnr")

	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	// Without going too far into gwacl's innards, this is roughly what
	// it needs to do in order to call GetContainerProperties.
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "GET")
}

func (*storageSuite) TestCreateContainerFailsIfContainerInaccessible(c *C) {
	azStorage, transport := makeFakeStorage("", "account")
	transport.AddExchange(makeResponse("", http.StatusInternalServerError), nil)

	err := azStorage.createContainer("cntnr")
	c.Assert(err, NotNil)

	// createContainer got an error when trying to query for an existing
	// container of the right name.  But it does not mistake that error for
	// "this container does not exist yet so go ahead and create it."
	// The proper response to the situation is to report the failure.
	c.Assert(err, ErrorMatches, ".*Internal Server Error.*")
}

func (*storageSuite) TestDeleteContainer(c *C) {
	azStorage, transport := makeFakeStorage("", "account")
	transport.AddExchange(makeResponse("", http.StatusAccepted), nil)

	err := azStorage.deleteContainer("cntnr")

	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	// Without going too far into gwacl's innards, this is roughly what
	// it needs to do in order to call GetContainerProperties.
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "DELETE")
}
<commit_msg>Set realistic endpoint for 'URL' storage test.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package azure

import (
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	. "launchpad.net/gocheck"
	"launchpad.net/gwacl"
	"launchpad.net/juju-core/errors"
)

type storageSuite struct {
	providerSuite
}

var _ = Suite(&storageSuite{})

func makeResponse(content string, status int) *http.Response {
	return &http.Response{
		Status:     fmt.Sprintf("%d", status),
		StatusCode: status,
		Body:       ioutil.NopCloser(strings.NewReader(content)),
	}
}

// MockingTransportExchange is a recording of a request and a response over
// HTTP.
type MockingTransportExchange struct {
	Request  *http.Request
	Response *http.Response
	Error    error
}

// MockingTransport is used as an http.Client.Transport for testing.  It
// records the sequence of requests, and returns a predetermined sequence of
// Responses and errors.
type MockingTransport struct {
	Exchanges     []*MockingTransportExchange
	ExchangeCount int
}

// MockingTransport implements the http.RoundTripper interface.
var _ http.RoundTripper = &MockingTransport{}

func (t *MockingTransport) AddExchange(response *http.Response, err error) {
	exchange := MockingTransportExchange{Response: response, Error: err}
	t.Exchanges = append(t.Exchanges, &exchange)
}

func (t *MockingTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	exchange := t.Exchanges[t.ExchangeCount]
	t.ExchangeCount++
	exchange.Request = req
	return exchange.Response, exchange.Error
}

// testStorageContext is a struct implementing the storageContext interface
// used in test.  It will return, via getContainer() and getStorageContext()
// the objects used at creation time.
type testStorageContext struct {
	container      string
	storageContext *gwacl.StorageContext
}

func (context *testStorageContext) getContainer() string {
	return context.container
}

func (context *testStorageContext) getStorageContext() (*gwacl.StorageContext, error) {
	return context.storageContext, nil
}

// makeFakeStorage creates a test azureStorage object that will talk to a
// fake HTTP server set up to always return preconfigured http.Response objects.
// The MockingTransport object can be used to check that the expected query has
// been issued to the test server.
func makeFakeStorage(container, account string) (azureStorage, *MockingTransport) {
	transport := &MockingTransport{}
	client := &http.Client{Transport: transport}
	storageContext := gwacl.NewTestStorageContext(client)
	storageContext.Account = account
	context := &testStorageContext{container: container, storageContext: storageContext}
	azStorage := azureStorage{storageContext: context}
	return azStorage, transport
}

// setStorageEndpoint sets a given Azure API endpoint on a given azureStorage.
func setStorageEndpoint(azStorage *azureStorage, endpoint gwacl.APIEndpoint) {
	// Ugly, because of the confusingly similar layers of nesting.
	testContext := azStorage.storageContext.(*testStorageContext)
	var gwaclContext *gwacl.StorageContext = testContext.storageContext
	gwaclContext.AzureEndpoint = endpoint
}

var blobListResponse = `
  <?xml version="1.0" encoding="utf-8"?>
  <EnumerationResults ContainerName="http://myaccount.blob.core.windows.net/mycontainer">
    <Prefix>prefix</Prefix>
    <Marker>marker</Marker>
    <MaxResults>maxresults</MaxResults>
    <Delimiter>delimiter</Delimiter>
    <Blobs>
      <Blob>
        <Name>prefix-1</Name>
        <Url>blob-url1</Url>
      </Blob>
      <Blob>
        <Name>prefix-2</Name>
        <Url>blob-url2</Url>
      </Blob>
    </Blobs>
    <NextMarker />
  </EnumerationResults>`

func (*storageSuite) TestList(c *C) {
	container := "container"
	response := makeResponse(blobListResponse, http.StatusOK)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)

	prefix := "prefix"
	names, err := azStorage.List(prefix)
	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	// The prefix has been passed down as a query parameter.
	c.Check(transport.Exchanges[0].Request.URL.Query()["prefix"], DeepEquals, []string{prefix})
	// The container name is used in the requested URL.
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, ".*"+container+".*")
	c.Check(names, DeepEquals, []string{"prefix-1", "prefix-2"})
}

func (*storageSuite) TestListWithNonexistentContainerReturnsNoFiles(c *C) {
	// If Azure returns a 404 it means the container doesn't exist. In this
	// case the provider should interpret this as "no files" and return nil.
	container := "container"
	response := makeResponse("", http.StatusNotFound)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)

	names, err := azStorage.List("prefix")
	c.Assert(err, IsNil)
	c.Assert(names, IsNil)
}

func (*storageSuite) TestGet(c *C) {
	blobContent := "test blob"
	container := "container"
	filename := "blobname"
	response := makeResponse(blobContent, http.StatusOK)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)

	reader, err := azStorage.Get(filename)
	c.Assert(err, IsNil)
	c.Assert(reader, NotNil)
	defer reader.Close()

	context, err := azStorage.getStorageContext()
	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, context.GetFileURL(container, filename)+"?.*")
	data, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(string(data), Equals, blobContent)
}

func (*storageSuite) TestGetReturnsNotFoundIf404(c *C) {
	container := "container"
	filename := "blobname"
	response := makeResponse("not found", http.StatusNotFound)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)
	_, err := azStorage.Get(filename)
	c.Assert(err, NotNil)
	c.Check(errors.IsNotFoundError(err), Equals, true)
}

func (*storageSuite) TestPut(c *C) {
	blobContent := "test blob"
	container := "container"
	filename := "blobname"
	azStorage, transport := makeFakeStorage(container, "account")
	// The create container call makes two exchanges.
	transport.AddExchange(makeResponse("", http.StatusNotFound), nil)
	transport.AddExchange(makeResponse("", http.StatusCreated), nil)
	putResponse := makeResponse("", http.StatusCreated)
	transport.AddExchange(putResponse, nil)
	transport.AddExchange(putResponse, nil)
	err := azStorage.Put(filename, strings.NewReader(blobContent), int64(len(blobContent)))
	c.Assert(err, IsNil)

	context, err := azStorage.getStorageContext()
	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 4)
	c.Check(transport.Exchanges[2].Request.URL.String(), Matches, context.GetFileURL(container, filename)+"?.*")
}

func (*storageSuite) TestRemove(c *C) {
	container := "container"
	filename := "blobname"
	response := makeResponse("", http.StatusAccepted)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)
	err := azStorage.Remove(filename)
	c.Assert(err, IsNil)

	context, err := azStorage.getStorageContext()
	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, context.GetFileURL(container, filename)+"?.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "DELETE")
}

func (*storageSuite) TestRemoveErrors(c *C) {
	container := "container"
	filename := "blobname"
	response := makeResponse("", http.StatusForbidden)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)
	err := azStorage.Remove(filename)
	c.Assert(err, NotNil)
}

func (*storageSuite) TestRemoveAll(c *C) {
	// When we ask gwacl to remove all blobs, it calls DeleteContainer.
	response := makeResponse("", http.StatusAccepted)
	storage, transport := makeFakeStorage("cntnr", "account")
	transport.AddExchange(response, nil)

	err := storage.RemoveAll()
	c.Assert(err, IsNil)

	_, err = storage.getStorageContext()
	c.Assert(err, IsNil)
	// Without going too far into gwacl's innards, this is roughly what
	// it needs to do in order to delete a container.
	c.Assert(transport.ExchangeCount, Equals, 1)
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "DELETE")
}

func (*storageSuite) TestRemoveNonExistentBlobSucceeds(c *C) {
	container := "container"
	filename := "blobname"
	response := makeResponse("", http.StatusNotFound)
	azStorage, transport := makeFakeStorage(container, "account")
	transport.AddExchange(response, nil)
	err := azStorage.Remove(filename)
	c.Assert(err, IsNil)
}

func (*storageSuite) TestURL(c *C) {
	container := "container"
	filename := "blobname"
	account := "account"
	azStorage, _ := makeFakeStorage(container, account)
	// Use a realistic service endpoint for this test, so that we can see
	// that we're really getting the expected kind of URL.
	setStorageEndpoint(&azStorage, gwacl.GetEndpoint("West US"))
	URL, err := azStorage.URL(filename)
	c.Assert(err, IsNil)
	parsedURL, err := url.Parse(URL)
	c.Assert(err, IsNil)
	c.Check(parsedURL.Host, Matches, fmt.Sprintf("%s.blob.core.windows.net", account))
	c.Check(parsedURL.Path, Matches, fmt.Sprintf("/%s/%s", container, filename))
	values, err := url.ParseQuery(parsedURL.RawQuery)
	c.Assert(err, IsNil)
	signature := values.Get("sig")
	// The query string contains a non-empty signature.
	c.Check(signature, Not(HasLen), 0)
	// The signature is base64-encoded.
	_, err = base64.StdEncoding.DecodeString(signature)
	c.Assert(err, IsNil)
}

func (*storageSuite) TestCreateContainerCreatesContainerIfDoesNotExist(c *C) {
	azStorage, transport := makeFakeStorage("", "account")
	transport.AddExchange(makeResponse("", http.StatusNotFound), nil)
	transport.AddExchange(makeResponse("", http.StatusCreated), nil)

	err := azStorage.createContainer("cntnr")

	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 2)
	// Without going too far into gwacl's innards, this is roughly what
	// it needs to do in order to call GetContainerProperties.
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "GET")

	// ... and for CreateContainer.
	c.Check(transport.Exchanges[1].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[1].Request.Method, Equals, "PUT")
}

func (*storageSuite) TestCreateContainerIsDoneIfContainerAlreadyExists(c *C) {
	container := ""
	azStorage, transport := makeFakeStorage(container, "account")
	header := make(http.Header)
	header.Add("Last-Modified", "last-modified")
	header.Add("ETag", "etag")
	header.Add("X-Ms-Lease-Status", "status")
	header.Add("X-Ms-Lease-State", "state")
	header.Add("X-Ms-Lease-Duration", "duration")
	response := makeResponse("", http.StatusOK)
	response.Header = header
	transport.AddExchange(response, nil)

	err := azStorage.createContainer("cntnr")

	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	// Without going too far into gwacl's innards, this is roughly what
	// it needs to do in order to call GetContainerProperties.
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "GET")
}

func (*storageSuite) TestCreateContainerFailsIfContainerInaccessible(c *C) {
	azStorage, transport := makeFakeStorage("", "account")
	transport.AddExchange(makeResponse("", http.StatusInternalServerError), nil)

	err := azStorage.createContainer("cntnr")
	c.Assert(err, NotNil)

	// createContainer got an error when trying to query for an existing
	// container of the right name.  But it does not mistake that error for
	// "this container does not exist yet so go ahead and create it."
	// The proper response to the situation is to report the failure.
	c.Assert(err, ErrorMatches, ".*Internal Server Error.*")
}

func (*storageSuite) TestDeleteContainer(c *C) {
	azStorage, transport := makeFakeStorage("", "account")
	transport.AddExchange(makeResponse("", http.StatusAccepted), nil)

	err := azStorage.deleteContainer("cntnr")

	c.Assert(err, IsNil)
	c.Assert(transport.ExchangeCount, Equals, 1)
	// Without going too far into gwacl's innards, this is roughly what
	// it needs to do in order to call GetContainerProperties.
	c.Check(transport.Exchanges[0].Request.URL.String(), Matches, "http.*/cntnr?.*restype=container.*")
	c.Check(transport.Exchanges[0].Request.Method, Equals, "DELETE")
}
<|endoftext|>


<commit_before>// Copyright 2018 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.
package phonenumbers

import (
	"testing"

	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"
)

func TestSetPhoneNumber(t *testing.T) {
	tc := libkb.SetupTest(t, "TestPhoneNumbers", 1)
	defer tc.Cleanup()

	me, err := kbtest.CreateAndSignupFakeUser("phon", tc.G)
	require.NoError(t, err)

	phoneNumber := keybase1.PhoneNumber("+14155552671")

	mctx := libkb.NewMetaContextForTest(tc)

	err = AddPhoneNumber(mctx, phoneNumber, keybase1.IdentityVisibility_PRIVATE)
	require.NoError(t, err)

	code, err := kbtest.GetPhoneVerificationCode(mctx, phoneNumber)
	require.NoError(t, err)
	t.Logf("Got verification code: %q", code)

	err = VerifyPhoneNumber(mctx, phoneNumber, code)
	require.NoError(t, err)
	err = SetVisibilityPhoneNumber(mctx, phoneNumber, keybase1.IdentityVisibility_PUBLIC)
	require.NoError(t, err)

	resp, err := GetPhoneNumbers(mctx)
	require.NoError(t, err)
	require.Len(t, resp, 1)
	require.Equal(t, phoneNumber, resp[0].PhoneNumber)
	require.True(t, resp[0].Verified)

	contactList := []keybase1.RawPhoneNumber{
		"+1-415-555-2671",
	}
	regionCodes := []keybase1.RegionCode{
		"us",
	}
	userRegionCode := keybase1.RegionCode("us")
	resolutions, err := BulkLookupPhoneNumbers(mctx, contactList, regionCodes, &userRegionCode)
	require.NoError(t, err)

	myUID := me.GetUID()
	expectedResolutions := []keybase1.PhoneNumberLookupResult{
		keybase1.PhoneNumberLookupResult{
			PhoneNumber:        "+1-415-555-2671",
			CoercedPhoneNumber: "+14155552671",
			Err:                nil,
			Uid:                &myUID,
		},
	}

	require.Equal(t, expectedResolutions, resolutions)

	err = DeletePhoneNumber(mctx, phoneNumber)
	require.NoError(t, err)

	resp, err = GetPhoneNumbers(mctx)
	require.NoError(t, err)
	require.Len(t, resp, 0)
}

func TestBadPhoneNumbers(t *testing.T) {
	tc := libkb.SetupTest(t, "TestPhoneNumbers", 1)
	defer tc.Cleanup()

	_, err := kbtest.CreateAndSignupFakeUser("phon", tc.G)
	require.NoError(t, err)

	mctx := libkb.NewMetaContextForTest(tc)
	require.Error(t, AddPhoneNumber(mctx, "14155552671", keybase1.IdentityVisibility_PUBLIC))
	require.Error(t, AddPhoneNumber(mctx, "014155552671", keybase1.IdentityVisibility_PUBLIC))
	require.Error(t, AddPhoneNumber(mctx, "784111222", keybase1.IdentityVisibility_PUBLIC))
}
<commit_msg>Fix TestSetPhoneNumber by generating random numbers (#16548)<commit_after>// Copyright 2018 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.
package phonenumbers

import (
	"fmt"
	"strings"
	"testing"

	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"
)

func TestSetPhoneNumber(t *testing.T) {
	tc := libkb.SetupTest(t, "TestPhoneNumbers", 1)
	defer tc.Cleanup()

	me, err := kbtest.CreateAndSignupFakeUser("phon", tc.G)
	require.NoError(t, err)

	// Generate a random phone number e.g. "14155552671".
	randomNumber := kbtest.GenerateTestPhoneNumber()
	// In strict format: "+14155552671".
	phoneNumber := keybase1.PhoneNumber("+" + randomNumber)
	// Create a representation likely to come from phone contact book: "+1-415-555-2671".
	phoneFormatted := keybase1.RawPhoneNumber(fmt.Sprintf("+%s-%s-%s-%s", randomNumber[0:1], randomNumber[1:4], randomNumber[4:7], randomNumber[7:11]))
	// Sanity check.
	require.EqualValues(t, phoneNumber, strings.Replace(string(phoneFormatted), "-", "", -1))

	t.Logf("Generated phone number: %q formatted as %q", phoneNumber, phoneFormatted)

	mctx := libkb.NewMetaContextForTest(tc)

	err = AddPhoneNumber(mctx, phoneNumber, keybase1.IdentityVisibility_PRIVATE)
	require.NoError(t, err)

	code, err := kbtest.GetPhoneVerificationCode(mctx, phoneNumber)
	require.NoError(t, err)
	t.Logf("Got verification code: %q", code)

	err = VerifyPhoneNumber(mctx, phoneNumber, code)
	require.NoError(t, err)
	err = SetVisibilityPhoneNumber(mctx, phoneNumber, keybase1.IdentityVisibility_PUBLIC)
	require.NoError(t, err)

	resp, err := GetPhoneNumbers(mctx)
	require.NoError(t, err)
	require.Len(t, resp, 1)
	require.Equal(t, phoneNumber, resp[0].PhoneNumber)
	require.True(t, resp[0].Verified)

	contactList := []keybase1.RawPhoneNumber{
		phoneFormatted,
	}
	regionCodes := []keybase1.RegionCode{
		"us",
	}
	userRegionCode := keybase1.RegionCode("us")
	resolutions, err := BulkLookupPhoneNumbers(mctx, contactList, regionCodes, &userRegionCode)
	require.NoError(t, err)

	myUID := me.GetUID()
	expectedResolutions := []keybase1.PhoneNumberLookupResult{
		keybase1.PhoneNumberLookupResult{
			PhoneNumber:        phoneFormatted,
			CoercedPhoneNumber: phoneNumber,
			Err:                nil,
			Uid:                &myUID,
		},
	}

	require.Equal(t, expectedResolutions, resolutions)

	err = DeletePhoneNumber(mctx, phoneNumber)
	require.NoError(t, err)

	resp, err = GetPhoneNumbers(mctx)
	require.NoError(t, err)
	require.Len(t, resp, 0)
}

func TestBadPhoneNumbers(t *testing.T) {
	tc := libkb.SetupTest(t, "TestPhoneNumbers", 1)
	defer tc.Cleanup()

	_, err := kbtest.CreateAndSignupFakeUser("phon", tc.G)
	require.NoError(t, err)

	mctx := libkb.NewMetaContextForTest(tc)
	require.Error(t, AddPhoneNumber(mctx, "14155552671", keybase1.IdentityVisibility_PUBLIC))
	require.Error(t, AddPhoneNumber(mctx, "014155552671", keybase1.IdentityVisibility_PUBLIC))
	require.Error(t, AddPhoneNumber(mctx, "784111222", keybase1.IdentityVisibility_PUBLIC))
}
<|endoftext|>


<commit_before>package main

import (
	"exp/inotify"
	"fmt"
	"io"
	"koding/tools/dnode"
	"koding/tools/kite"
	"koding/tools/log"
	"koding/virt"
	"os"
	"path"
	"regexp"
	"strconv"
	"sync"
	"time"
)

var watcher *inotify.Watcher
var watchMap = make(map[string][]*Watch)
var watchMutex sync.Mutex

type Watch struct {
	VOS      *virt.VOS
	Path     string
	Callback dnode.Callback
}

type FileEntry struct {
	Name     string      `json:"name"`
	IsDir    bool        `json:"isDir"`
	Size     int64       `json:"size"`
	Mode     os.FileMode `json:"mode"`
	Time     time.Time   `json:"time"`
	IsBroken bool        `json:"isBroken"`
}

func init() {
	var err error
	watcher, err = inotify.NewWatcher()
	if err != nil {
		panic(err)
	}

	go func() {
		for ev := range watcher.Event {
			if (ev.Mask & (inotify.IN_CREATE | inotify.IN_MODIFY | inotify.IN_MOVED_TO)) != 0 {
				info, err := os.Lstat(ev.Name)
				if err != nil {
					log.Warn("Watcher error", err)
					continue
				}
				watchMutex.Lock()
				for _, watch := range watchMap[path.Dir(ev.Name)] {
					watch.Callback(map[string]interface{}{
						"event": "added",
						"file":  makeFileEntry(watch.VOS, watch.Path, info),
					})
				}
				watchMutex.Unlock()
				continue
			}
			if (ev.Mask & (inotify.IN_DELETE | inotify.IN_MOVED_FROM)) != 0 {
				watchMutex.Lock()
				for _, watch := range watchMap[path.Dir(ev.Name)] {
					watch.Callback(map[string]interface{}{
						"event": "removed",
						"file":  FileEntry{Name: path.Base(ev.Name)},
					})
				}
				watchMutex.Unlock()
				continue
			}
		}
	}()
	go func() {
		for err := range watcher.Error {
			log.Warn("Watcher error", err)
		}
	}()
}

func registerFileSystemMethods(k *kite.Kite) {
	k.Handle("fs.readDirectory", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path     string
			OnChange dnode.Callback
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string], onChange: [function] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)
		response := make(map[string]interface{})

		if params.OnChange != nil {
			watchedPath, err := vos.AddWatch(watcher, params.Path, inotify.IN_CREATE|inotify.IN_DELETE|inotify.IN_MODIFY|inotify.IN_MOVE)
			if err != nil {
				return nil, err
			}

			watchMutex.Lock()
			defer watchMutex.Unlock()

			watch := &Watch{vos, watchedPath, params.OnChange}
			watchMap[watchedPath] = append(watchMap[watchedPath], watch)
			session.OnDisconnect(func() { watch.Close() })
			response["stopWatching"] = func() { watch.Close() }
		}

		dir, err := vos.Open(params.Path)
		defer dir.Close()
		if err != nil {
			return nil, err
		}

		infos, err := dir.Readdir(0)
		if err != nil {
			return nil, err
		}

		files := make([]FileEntry, len(infos))
		for i, info := range infos {
			files[i] = makeFileEntry(vos, params.Path, info)
		}
		response["files"] = files

		return response, nil
	})

	k.Handle("fs.readFile", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path string
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		file, err := vos.Open(params.Path)
		if err != nil {
			return nil, err
		}
		defer file.Close()

		fi, err := file.Stat()
		if err != nil {
			return nil, err
		}

		if fi.Size() > 10*1024*1024 {
			return nil, fmt.Errorf("File larger than 10MiB.")
		}

		buf := make([]byte, fi.Size())
		if _, err := io.ReadFull(file, buf); err != nil {
			return nil, err
		}

		return map[string]interface{}{"content": buf}, nil
	})

	k.Handle("fs.writeFile", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path           string
			Content        []byte
			DoNotOverwrite bool
		}
		if args.Unmarshal(&params) != nil || params.Path == "" || params.Content == nil {
			return nil, &kite.ArgumentError{Expected: "{ path: [string], content: [base64], doNotOverwrite: [bool] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		flags := os.O_RDWR | os.O_CREATE | os.O_TRUNC
		if params.DoNotOverwrite {
			flags |= os.O_EXCL
		}
		file, err := vos.OpenFile(params.Path, flags, 0666)
		if err != nil {
			return nil, err
		}
		defer file.Close()

		return file.Write(params.Content)
	})

	suffixRegexp, err := regexp.Compile(`.((_\d+)?)(\.\w*)?$`)
	if err != nil {
		panic(err)
	}

	k.Handle("fs.ensureNonexistentPath", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path string
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		name := params.Path
		index := 1
		for {
			_, err := vos.Stat(name)
			if err != nil {
				if os.IsNotExist(err) {
					break
				}
				return nil, err
			}

			loc := suffixRegexp.FindStringSubmatchIndex(name)
			name = name[:loc[2]] + "_" + strconv.Itoa(index) + name[loc[3]:]
			index++
		}

		return name, nil
	})

	k.Handle("fs.getInfo", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path string
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		fi, err := vos.Stat(params.Path)
		if err != nil {
			if os.IsNotExist(err) {
				return nil, nil
			}
			return nil, err
		}

		return makeFileEntry(vos, path.Dir(params.Path), fi), nil
	})

	k.Handle("fs.setPermissions", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path      string
			Mode      os.FileMode
			Recursive bool
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string], mode: [integer], recursive: [bool] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		doChange := func(name string) error {
			if err := vos.Chmod(name, params.Mode); err != nil {
				return err
			}
			if !params.Recursive {
				return nil
			}
			fi, err := vos.Stat(name)
			if err != nil {
				return err
			}
			if !fi.IsDir() {
				return nil
			}
			dir, err := vos.Open(name)
			if err != nil {
				return err
			}
			defer dir.Close()
			entries, err := dir.Readdirnames(0)
			if err != nil {
				return err
			}
			var firstErr error
			for entry := range entries {
				err := doChange(name + "/" + entry)
				if err != nil && firstErr == nil {
					firstErr = err
				}
			}
			return firstErr
		}
		if err := doChange(params.Path); err != nil {
			return nil, err
		}

		return true, nil
	})

	k.Handle("fs.remove", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path      string
			Recursive bool
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string], recursive: [bool] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		if params.Recursive {
			if err := vos.RemoveAll(params.Path); err != nil {
				return nil, err
			}
			return true, nil
		}

		if err := vos.Remove(params.Path); err != nil {
			return nil, err
		}
		return true, nil
	})

	k.Handle("fs.rename", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			OldPath string
			NewPath string
		}
		if args.Unmarshal(&params) != nil || params.OldPath == "" || params.NewPath == "" {
			return nil, &kite.ArgumentError{Expected: "{ oldPath: [string], newPath: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		if err := vos.Rename(params.OldPath, params.NewPath); err != nil {
			return nil, err
		}

		return true, nil
	})

	k.Handle("fs.createDirectory", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path string
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		if err := vos.Mkdir(params.Path, 0755); err != nil {
			return nil, err
		}

		return true, nil
	})
}

func (watch *Watch) Close() error {
	watchMutex.Lock()
	defer watchMutex.Unlock()

	watches := watchMap[watch.Path]
	for i, w := range watches {
		if w == watch {
			watches[i] = watches[len(watches)-1]
			watches = watches[:len(watches)-1]
			break
		}
	}

	watchMap[watch.Path] = watches

	if len(watches) == 0 {
		return watcher.RemoveWatch(watch.Path)
	}

	return nil
}

func makeFileEntry(vos *virt.VOS, dir string, fi os.FileInfo) FileEntry {
	entry := FileEntry{
		Name:  fi.Name(),
		IsDir: fi.IsDir(),
		Size:  fi.Size(),
		Mode:  fi.Mode(),
		Time:  fi.ModTime(),
	}

	if fi.Mode()&os.ModeSymlink != 0 {
		symlinkInfo, err := vos.Stat(dir + "/" + fi.Name())
		if err != nil {
			entry.IsBroken = true
			return entry
		}
		entry.IsDir = symlinkInfo.IsDir()
		entry.Size = symlinkInfo.Size()
		entry.Mode = symlinkInfo.Mode()
		entry.Time = symlinkInfo.ModTime()
	}

	return entry
}
<commit_msg>Fixes for previous commit.<commit_after>package main

import (
	"exp/inotify"
	"fmt"
	"io"
	"koding/tools/dnode"
	"koding/tools/kite"
	"koding/tools/log"
	"koding/virt"
	"os"
	"path"
	"regexp"
	"strconv"
	"sync"
	"time"
)

var watcher *inotify.Watcher
var watchMap = make(map[string][]*Watch)
var watchMutex sync.Mutex

type Watch struct {
	VOS      *virt.VOS
	Path     string
	Callback dnode.Callback
}

type FileEntry struct {
	Name     string      `json:"name"`
	IsDir    bool        `json:"isDir"`
	Size     int64       `json:"size"`
	Mode     os.FileMode `json:"mode"`
	Time     time.Time   `json:"time"`
	IsBroken bool        `json:"isBroken"`
}

func init() {
	var err error
	watcher, err = inotify.NewWatcher()
	if err != nil {
		panic(err)
	}

	go func() {
		for ev := range watcher.Event {
			if (ev.Mask & (inotify.IN_CREATE | inotify.IN_MODIFY | inotify.IN_MOVED_TO)) != 0 {
				info, err := os.Lstat(ev.Name)
				if err != nil {
					log.Warn("Watcher error", err)
					continue
				}
				watchMutex.Lock()
				for _, watch := range watchMap[path.Dir(ev.Name)] {
					watch.Callback(map[string]interface{}{
						"event": "added",
						"file":  makeFileEntry(watch.VOS, watch.Path, info),
					})
				}
				watchMutex.Unlock()
				continue
			}
			if (ev.Mask & (inotify.IN_DELETE | inotify.IN_MOVED_FROM)) != 0 {
				watchMutex.Lock()
				for _, watch := range watchMap[path.Dir(ev.Name)] {
					watch.Callback(map[string]interface{}{
						"event": "removed",
						"file":  FileEntry{Name: path.Base(ev.Name)},
					})
				}
				watchMutex.Unlock()
				continue
			}
		}
	}()
	go func() {
		for err := range watcher.Error {
			log.Warn("Watcher error", err)
		}
	}()
}

func registerFileSystemMethods(k *kite.Kite) {
	k.Handle("fs.readDirectory", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path     string
			OnChange dnode.Callback
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string], onChange: [function] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)
		response := make(map[string]interface{})

		if params.OnChange != nil {
			watchedPath, err := vos.AddWatch(watcher, params.Path, inotify.IN_CREATE|inotify.IN_DELETE|inotify.IN_MODIFY|inotify.IN_MOVE)
			if err != nil {
				return nil, err
			}

			watchMutex.Lock()
			defer watchMutex.Unlock()

			watch := &Watch{vos, watchedPath, params.OnChange}
			watchMap[watchedPath] = append(watchMap[watchedPath], watch)
			session.OnDisconnect(func() { watch.Close() })
			response["stopWatching"] = func() { watch.Close() }
		}

		dir, err := vos.Open(params.Path)
		defer dir.Close()
		if err != nil {
			return nil, err
		}

		infos, err := dir.Readdir(0)
		if err != nil {
			return nil, err
		}

		files := make([]FileEntry, len(infos))
		for i, info := range infos {
			files[i] = makeFileEntry(vos, params.Path, info)
		}
		response["files"] = files

		return response, nil
	})

	k.Handle("fs.readFile", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path string
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		file, err := vos.Open(params.Path)
		if err != nil {
			return nil, err
		}
		defer file.Close()

		fi, err := file.Stat()
		if err != nil {
			return nil, err
		}

		if fi.Size() > 10*1024*1024 {
			return nil, fmt.Errorf("File larger than 10MiB.")
		}

		buf := make([]byte, fi.Size())
		if _, err := io.ReadFull(file, buf); err != nil {
			return nil, err
		}

		return map[string]interface{}{"content": buf}, nil
	})

	k.Handle("fs.writeFile", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path           string
			Content        []byte
			DoNotOverwrite bool
		}
		if args.Unmarshal(&params) != nil || params.Path == "" || params.Content == nil {
			return nil, &kite.ArgumentError{Expected: "{ path: [string], content: [base64], doNotOverwrite: [bool] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		flags := os.O_RDWR | os.O_CREATE | os.O_TRUNC
		if params.DoNotOverwrite {
			flags |= os.O_EXCL
		}
		file, err := vos.OpenFile(params.Path, flags, 0666)
		if err != nil {
			return nil, err
		}
		defer file.Close()

		return file.Write(params.Content)
	})

	suffixRegexp, err := regexp.Compile(`.((_\d+)?)(\.\w*)?$`)
	if err != nil {
		panic(err)
	}

	k.Handle("fs.ensureNonexistentPath", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path string
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		name := params.Path
		index := 1
		for {
			_, err := vos.Stat(name)
			if err != nil {
				if os.IsNotExist(err) {
					break
				}
				return nil, err
			}

			loc := suffixRegexp.FindStringSubmatchIndex(name)
			name = name[:loc[2]] + "_" + strconv.Itoa(index) + name[loc[3]:]
			index++
		}

		return name, nil
	})

	k.Handle("fs.getInfo", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path string
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		fi, err := vos.Stat(params.Path)
		if err != nil {
			if os.IsNotExist(err) {
				return nil, nil
			}
			return nil, err
		}

		return makeFileEntry(vos, path.Dir(params.Path), fi), nil
	})

	k.Handle("fs.setPermissions", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path      string
			Mode      os.FileMode
			Recursive bool
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string], mode: [integer], recursive: [bool] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		var doChange func(name string) error
		doChange = func(name string) error {
			if err := vos.Chmod(name, params.Mode); err != nil {
				return err
			}
			if !params.Recursive {
				return nil
			}
			fi, err := vos.Stat(name)
			if err != nil {
				return err
			}
			if !fi.IsDir() {
				return nil
			}
			dir, err := vos.Open(name)
			if err != nil {
				return err
			}
			defer dir.Close()
			entries, err := dir.Readdirnames(0)
			if err != nil {
				return err
			}
			var firstErr error
			for _, entry := range entries {
				err := doChange(name + "/" + entry)
				if err != nil && firstErr == nil {
					firstErr = err
				}
			}
			return firstErr
		}
		if err := doChange(params.Path); err != nil {
			return nil, err
		}

		return true, nil
	})

	k.Handle("fs.remove", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path      string
			Recursive bool
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string], recursive: [bool] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		if params.Recursive {
			if err := vos.RemoveAll(params.Path); err != nil {
				return nil, err
			}
			return true, nil
		}

		if err := vos.Remove(params.Path); err != nil {
			return nil, err
		}
		return true, nil
	})

	k.Handle("fs.rename", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			OldPath string
			NewPath string
		}
		if args.Unmarshal(&params) != nil || params.OldPath == "" || params.NewPath == "" {
			return nil, &kite.ArgumentError{Expected: "{ oldPath: [string], newPath: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		if err := vos.Rename(params.OldPath, params.NewPath); err != nil {
			return nil, err
		}

		return true, nil
	})

	k.Handle("fs.createDirectory", false, func(args *dnode.Partial, session *kite.Session) (interface{}, error) {
		var params struct {
			Path string
		}
		if args.Unmarshal(&params) != nil || params.Path == "" {
			return nil, &kite.ArgumentError{Expected: "{ path: [string] }"}
		}

		user, vm := findSession(session)
		vos := vm.OS(user)

		if err := vos.Mkdir(params.Path, 0755); err != nil {
			return nil, err
		}

		return true, nil
	})
}

func (watch *Watch) Close() error {
	watchMutex.Lock()
	defer watchMutex.Unlock()

	watches := watchMap[watch.Path]
	for i, w := range watches {
		if w == watch {
			watches[i] = watches[len(watches)-1]
			watches = watches[:len(watches)-1]
			break
		}
	}

	watchMap[watch.Path] = watches

	if len(watches) == 0 {
		return watcher.RemoveWatch(watch.Path)
	}

	return nil
}

func makeFileEntry(vos *virt.VOS, dir string, fi os.FileInfo) FileEntry {
	entry := FileEntry{
		Name:  fi.Name(),
		IsDir: fi.IsDir(),
		Size:  fi.Size(),
		Mode:  fi.Mode(),
		Time:  fi.ModTime(),
	}

	if fi.Mode()&os.ModeSymlink != 0 {
		symlinkInfo, err := vos.Stat(dir + "/" + fi.Name())
		if err != nil {
			entry.IsBroken = true
			return entry
		}
		entry.IsDir = symlinkInfo.IsDir()
		entry.Size = symlinkInfo.Size()
		entry.Mode = symlinkInfo.Mode()
		entry.Time = symlinkInfo.ModTime()
	}

	return entry
}
<|endoftext|>


<commit_before>package uidmap

import (
	"testing"
	"time"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/clockwork"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

const tTracy = keybase1.UID("eb72f49f2dde6429e5d78003dae0c919")
const tAlice = keybase1.UID("295a7eea607af32040647123732bc819")

// Larger than 1 nanosecond which would skip request altogether. Use to
// "simulate" request that didn't make it back in time.
const strictNetworkBudget = 2 * time.Nanosecond

func TestServiceMapLookupKnown(t *testing.T) {
	t.Skip()

	tc := libkb.SetupTest(t, "TestLookup", 1)
	defer tc.Cleanup()

	fakeClock := clockwork.NewFakeClockAt(time.Now())
	tc.G.SetClock(fakeClock)

	now := keybase1.ToTime(fakeClock.Now())

	serviceMapper := NewServiceSummaryMap(10)
	uids := []keybase1.UID{tKB, tAlice, tTracy}
	const zeroDuration = time.Duration(0)
	pkgs := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
		zeroDuration /* freshness */, zeroDuration /* networkBudget */)

	require.Len(t, pkgs, 3)
	require.Contains(t, pkgs, tKB)
	require.Contains(t, pkgs, tAlice)
	require.Contains(t, pkgs, tTracy)
	for _, v := range pkgs {
		require.True(t, v.CachedAt >= now)
		require.NotNil(t, v.ServiceMap)
	}

	// Exact maps depend on remote_identities on the test server.
	require.Equal(t, "gbrltest", pkgs[tKB].ServiceMap["twitter"])
	require.Equal(t, "tacovontaco", pkgs[tAlice].ServiceMap["twitter"])
	require.Equal(t, "tacoplusplus", pkgs[tTracy].ServiceMap["github"])
	require.Equal(t, "t_tracy", pkgs[tTracy].ServiceMap["rooter"])
	require.Equal(t, "tacovontaco", pkgs[tTracy].ServiceMap["twitter"])

	{
		// Query again with very strict network budget hoping to hit cache.
		pkgs2 := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
			zeroDuration /* freshness */, strictNetworkBudget /* networkBudget */)
		require.Equal(t, pkgs, pkgs2)
	}

	{
		// Same, but advance fake clock and provide `freshness` argument. We
		// should fail to get data.
		fakeClock.Advance(24 * time.Hour)

		pkgs2 := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
			12*time.Hour, /* freshness */
			strictNetworkBudget /* networkBudget */)
		require.Len(t, pkgs2, 0)
	}

	{
		// Similar, but with DisallowNetworkBudget which should skip request completely.
		pkgs2 := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
			12*time.Hour, /* freshness */
			DisallowNetworkBudget /* networkBudget */)
		require.Len(t, pkgs2, 0)
	}
}

func TestServiceMapLookupEmpty(t *testing.T) {
	tc := libkb.SetupTest(t, "TestLookup", 1)
	defer tc.Cleanup()

	now := keybase1.ToTime(time.Now())
	serviceMapper := NewServiceSummaryMap(10)

	const tFrank = keybase1.UID("359c7644857203be38bfd3bf79bf1819")
	uids := []keybase1.UID{tFrank}
	const zeroDuration = time.Duration(0)
	pkgs := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
		zeroDuration /* freshness */, zeroDuration /* networkBudget */)

	// t_frank has no services, expecting to see t_frank in result map but with
	// nil ServiceMap field.
	require.Len(t, pkgs, 1)
	require.Contains(t, pkgs, tFrank)
	require.Nil(t, pkgs[tFrank].ServiceMap)
	require.True(t, pkgs[tFrank].CachedAt >= now)

	{
		// Query again with very strict network budget hoping to hit cache.
		pkgs2 := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
			zeroDuration /* freshness */, strictNetworkBudget /* networkBudget */)
		require.Equal(t, pkgs, pkgs2)
	}
}
<commit_msg>Use API mocks in servicemap_test.go (#22424)<commit_after>package uidmap

import (
	"errors"
	"testing"
	"time"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/clockwork"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

type timeoutAPIMock struct {
	libkb.API
	callCount int
}

func (n *timeoutAPIMock) PostDecodeCtx(context.Context, libkb.APIArg, libkb.APIResponseWrapper) error {
	n.callCount++
	return libkb.APINetError{Err: errors.New("timeoutAPIMock")}
}

const tTracy = keybase1.UID("eb72f49f2dde6429e5d78003dae0c919")
const tAlice = keybase1.UID("295a7eea607af32040647123732bc819")

func TestServiceMapLookupKnown(t *testing.T) {
	tc := libkb.SetupTest(t, "TestLookup", 1)
	defer tc.Cleanup()

	fakeClock := clockwork.NewFakeClockAt(time.Now())
	tc.G.SetClock(fakeClock)

	now := keybase1.ToTime(fakeClock.Now())

	serviceMapper := NewServiceSummaryMap(10)
	uids := []keybase1.UID{tKB, tAlice, tTracy}
	const zeroDuration = time.Duration(0)
	pkgs := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
		zeroDuration /* freshness */, DefaultNetworkBudget /* networkBudget */)

	require.Len(t, pkgs, 3)
	require.Contains(t, pkgs, tKB)
	require.Contains(t, pkgs, tAlice)
	require.Contains(t, pkgs, tTracy)
	for _, v := range pkgs {
		require.True(t, v.CachedAt >= now)
		require.NotNil(t, v.ServiceMap)
	}

	// Exact maps depend on remote_identities on the test server.
	require.Equal(t, "gbrltest", pkgs[tKB].ServiceMap["twitter"])
	require.Equal(t, "tacovontaco", pkgs[tAlice].ServiceMap["twitter"])
	require.Equal(t, "tacoplusplus", pkgs[tTracy].ServiceMap["github"])
	require.Equal(t, "t_tracy", pkgs[tTracy].ServiceMap["rooter"])
	require.Equal(t, "tacovontaco", pkgs[tTracy].ServiceMap["twitter"])

	timeoutAPI := &timeoutAPIMock{}
	tc.G.API = timeoutAPI

	// Doesn't matter for the API itself because we are mocking it, but make
	// sure serviceMapper does not do anything funky when budget is anything
	// other than default.
	networkBudget := 1 * time.Second

	{
		pkgs2 := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
			zeroDuration /* freshness */, networkBudget)
		require.Equal(t, pkgs, pkgs2)
		// Should not have to call API with freshness set to "always fresh".
		require.Equal(t, 0, timeoutAPI.callCount)
	}

	{
		// Same, but advance fake clock and provide `freshness` argument. We
		// should fail to get data with always-timeouting API.
		fakeClock.Advance(24 * time.Hour)

		pkgs2 := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
			12*time.Hour /* freshness */, networkBudget)
		require.Len(t, pkgs2, 0)
		require.Equal(t, 1, timeoutAPI.callCount)
	}

	{
		// Similar, but with DisallowNetworkBudget which should skip request completely.
		pkgs2 := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
			12*time.Hour /* freshness */, DisallowNetworkBudget /* networkBudget */)
		require.Len(t, pkgs2, 0)
		require.Equal(t, 1, timeoutAPI.callCount) // same count as after previous call
	}
}

func TestServiceMapLookupEmpty(t *testing.T) {
	tc := libkb.SetupTest(t, "TestLookup", 1)
	defer tc.Cleanup()

	now := keybase1.ToTime(time.Now())
	serviceMapper := NewServiceSummaryMap(10)

	const tFrank = keybase1.UID("359c7644857203be38bfd3bf79bf1819")
	uids := []keybase1.UID{tFrank}
	const zeroDuration = time.Duration(0)
	pkgs := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
		zeroDuration /* freshness */, zeroDuration /* networkBudget */)

	// t_frank has no services, expecting to see t_frank in result map but with
	// nil ServiceMap field.
	require.Len(t, pkgs, 1)
	require.Contains(t, pkgs, tFrank)
	require.Nil(t, pkgs[tFrank].ServiceMap)
	require.True(t, pkgs[tFrank].CachedAt >= now)

	timeoutAPI := &timeoutAPIMock{}
	tc.G.API = timeoutAPI

	{
		// Query again with very strict network budget hoping to hit cache.
		pkgs2 := serviceMapper.MapUIDsToServiceSummaries(context.TODO(), tc.G, uids,
			zeroDuration /* freshness */, DefaultNetworkBudget /* networkBudget */)
		require.Equal(t, pkgs, pkgs2)
		// should not hit the API due to freshness parameter.
		require.Equal(t, 0, timeoutAPI.callCount)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package vitessdriver

import (
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/topo/topoproto"
	"github.com/youtube/vitess/go/vt/vtgate/vtgateconn"
	"golang.org/x/net/context"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

func init() {
	sql.Register("vitess", drv{})
}

// TODO(mberlin): Add helper methods.

type drv struct {
}

// Open must be called with a JSON string that looks like this:
//
//   {"protocol": "gorpc", "address": "localhost:1111", "tablet_type": "master", "timeout": 1000000000}
//
// protocol specifies the rpc protocol to use.
// address specifies the address for the VTGate to connect to.
// tablet_type represents the consistency level of your operations.
// For example "replica" means eventually consistent reads, while
// "master" supports transactions and gives you read-after-write consistency.
// timeout is specified in nanoseconds. It applies for all operations.
//
// If you want to execute queries which are not supported by vtgate v3, you can
// run queries against a specific keyspace and shard.
// Therefore, add the fields "keyspace" and "shard" to the JSON string. Example:
//
//   {"protocol": "gorpc", "address": "localhost:1111", "keyspace": "ks1", "shard": "0", "tablet_type": "master", "timeout": 1000000000}
//
// Note that this function will always create a connection to vtgate i.e. there
// is no need to call DB.Ping() to verify the connection.
func (d drv) Open(name string) (driver.Conn, error) {
	c := &conn{TabletType: "master"}
	err := json.Unmarshal([]byte(name), c)
	if err != nil {
		return nil, err
	}
	if (c.Keyspace != "" && c.Shard == "") || (c.Keyspace == "" && c.Shard != "") {
		return nil, fmt.Errorf("Always set both keyspace and shard or leave both empty. keyspace: %v shard: %v", c.Keyspace, c.Shard)
	}
	if c.useExecuteShards() {
		log.Infof("Sending queries only to keyspace/shard: %v/%v", c.Keyspace, c.Shard)
	}
	c.tabletType, err = topoproto.ParseTabletType(c.TabletType)
	if err != nil {
		return nil, err
	}
	err = c.dial()
	if err != nil {
		return nil, err
	}
	return c, nil
}

type conn struct {
	Protocol string
	Address  string
	// Keyspace of a specific keyspace/shard to target. Disables vtgate v3.
	// If Keyspace and Shard are not empty, vtgate v2 instead of v3 will be used
	// and all requests will be sent only to that particular shard.
	// This functionality is meant for initial migrations from MySQL/MariaDB to Vitess.
	Keyspace string
	// Shard of a specific keyspace/shard to target. Disables vtgate v3.
	Shard      string
	TabletType string `json:"tablet_type"`
	Streaming  bool
	Timeout    time.Duration

	tabletType topodatapb.TabletType
	vtgateConn *vtgateconn.VTGateConn
	tx         *vtgateconn.VTGateTx
}

func (c *conn) dial() error {
	var err error
	if c.Protocol == "" {
		c.vtgateConn, err = vtgateconn.Dial(context.Background(), c.Address, c.Timeout)
	} else {
		c.vtgateConn, err = vtgateconn.DialProtocol(context.Background(), c.Protocol, c.Address, c.Timeout)
	}
	return err
}

func (c *conn) useExecuteShards() bool {
	return c.Keyspace != "" && c.Shard != ""
}

func (c *conn) Prepare(query string) (driver.Stmt, error) {
	return &stmt{c: c, query: query}, nil
}

func (c *conn) Close() error {
	c.vtgateConn.Close()
	return nil
}

func (c *conn) Begin() (driver.Tx, error) {
	if c.Streaming {
		return nil, errors.New("transaction not allowed for streaming connection")
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.Timeout)
	defer cancel()
	tx, err := c.vtgateConn.Begin(ctx)
	if err != nil {
		return nil, err
	}
	c.tx = tx
	return c, nil
}

func (c *conn) Commit() error {
	if c.tx == nil {
		return errors.New("commit: not in transaction")
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.Timeout)
	defer func() {
		cancel()
		c.tx = nil
	}()
	return c.tx.Commit(ctx)
}

func (c *conn) Rollback() error {
	if c.tx == nil {
		return nil
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.Timeout)
	defer func() {
		cancel()
		c.tx = nil
	}()
	return c.tx.Rollback(ctx)
}

type stmt struct {
	c     *conn
	query string
}

func (s *stmt) Close() error {
	return nil
}

func (s *stmt) NumInput() int {
	return -1
}

func (s *stmt) Exec(args []driver.Value) (driver.Result, error) {
	ctx, cancel := context.WithTimeout(context.Background(), s.c.Timeout)
	defer cancel()

	if s.c.Streaming {
		return nil, errors.New("Exec not allowed for streaming connections")
	}

	qr, err := s.executeVitess(ctx, args)
	if err != nil {
		return nil, err
	}
	return result{int64(qr.InsertID), int64(qr.RowsAffected)}, nil
}

func (s *stmt) Query(args []driver.Value) (driver.Rows, error) {
	ctx, cancel := context.WithTimeout(context.Background(), s.c.Timeout)

	if s.c.Streaming {
		var qrc <-chan *sqltypes.Result
		var errFunc vtgateconn.ErrFunc
		var err error
		if s.c.useExecuteShards() {
			qrc, errFunc, err = s.c.vtgateConn.StreamExecuteShards(ctx, s.query, s.c.Keyspace, []string{s.c.Shard}, makeBindVars(args), s.c.tabletType)
		} else {
			qrc, errFunc, err = s.c.vtgateConn.StreamExecute(ctx, s.query, makeBindVars(args), s.c.tabletType)
		}
		if err != nil {
			return nil, err
		}
		return newStreamingRows(qrc, errFunc, cancel), nil
	}
	// Do not cancel in case of a streaming query. It will do it itself.
	defer cancel()

	qr, err := s.executeVitess(ctx, args)
	if err != nil {
		return nil, err
	}
	return newRows(qr), nil
}

func (s *stmt) executeVitess(ctx context.Context, args []driver.Value) (*sqltypes.Result, error) {
	if s.c.tx != nil {
		if s.c.useExecuteShards() {
			return s.c.tx.ExecuteShards(ctx, s.query, s.c.Keyspace, []string{s.c.Shard}, makeBindVars(args), s.c.tabletType, false /* notInTransaction */)
		}
		return s.c.tx.Execute(ctx, s.query, makeBindVars(args), s.c.tabletType, false /* notInTransaction */)
	}

	// Non-transactional case.
	if s.c.useExecuteShards() {
		return s.c.vtgateConn.ExecuteShards(ctx, s.query, s.c.Keyspace, []string{s.c.Shard}, makeBindVars(args), s.c.tabletType)
	}
	return s.c.vtgateConn.Execute(ctx, s.query, makeBindVars(args), s.c.tabletType)
}

func makeBindVars(args []driver.Value) map[string]interface{} {
	if len(args) == 0 {
		return map[string]interface{}{}
	}
	bv := make(map[string]interface{}, len(args))
	for i, v := range args {
		bv[fmt.Sprintf("v%d", i+1)] = v
	}
	return bv
}

type result struct {
	insertid, rowsaffected int64
}

func (r result) LastInsertId() (int64, error) {
	return r.insertid, nil
}

func (r result) RowsAffected() (int64, error) {
	return r.rowsaffected, nil
}
<commit_msg>vitessdriver: Factor out configuration into own struct.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package vitessdriver

import (
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/topo/topoproto"
	"github.com/youtube/vitess/go/vt/vtgate/vtgateconn"
	"golang.org/x/net/context"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

func init() {
	sql.Register("vitess", drv{})
}

// TODO(mberlin): Add helper methods.

type drv struct {
}

// Open must be called with a JSON string that looks like this:
//
//   {"protocol": "gorpc", "address": "localhost:1111", "tablet_type": "master", "timeout": 1000000000}
//
// protocol specifies the rpc protocol to use.
// address specifies the address for the VTGate to connect to.
// tablet_type represents the consistency level of your operations.
// For example "replica" means eventually consistent reads, while
// "master" supports transactions and gives you read-after-write consistency.
// timeout is specified in nanoseconds. It applies for all operations.
//
// If you want to execute queries which are not supported by vtgate v3, you can
// run queries against a specific keyspace and shard.
// Therefore, add the fields "keyspace" and "shard" to the JSON string. Example:
//
//   {"protocol": "gorpc", "address": "localhost:1111", "keyspace": "ks1", "shard": "0", "tablet_type": "master", "timeout": 1000000000}
//
// Note that this function will always create a connection to vtgate i.e. there
// is no need to call DB.Ping() to verify the connection.
func (d drv) Open(name string) (driver.Conn, error) {
	c := &conn{Configuration: newDefaultConfiguration()}
	err := json.Unmarshal([]byte(name), c)
	if err != nil {
		return nil, err
	}
	if (c.Keyspace != "" && c.Shard == "") || (c.Keyspace == "" && c.Shard != "") {
		return nil, fmt.Errorf("Always set both keyspace and shard or leave both empty. keyspace: %v shard: %v", c.Keyspace, c.Shard)
	}
	if c.useExecuteShards() {
		log.Infof("Sending queries only to keyspace/shard: %v/%v", c.Keyspace, c.Shard)
	}
	c.tabletTypeProto, err = topoproto.ParseTabletType(c.TabletType)
	if err != nil {
		return nil, err
	}
	err = c.dial()
	if err != nil {
		return nil, err
	}
	return c, nil
}

// Configuration holds all Vitess driver settings.
//
// Fields with documented default values do not have to be set explicitly.
type Configuration struct {
	// Protocol is the name of the vtgate RPC client implementation.
	// Note: In open-source "grpc" is the recommended implementation.
	//
	// Default: "grpc"
	Protocol string

	// Address must point to a vtgate instance.
	//
	// Format: hostname:port
	Address string

	// Keyspace of a specific keyspace and shard to target. Disables vtgate v3.
	//
	// If Keyspace and Shard are not empty, vtgate v2 instead of v3 will be used
	// and all requests will be sent only to that particular shard.
	// This functionality is meant for initial migrations from MySQL/MariaDB to Vitess.
	Keyspace string
	// Shard of a specific keyspace and shard to target. Disables vtgate v3.
	Shard string

	// TabletType is the type of tablet you want to access and affects the
	// freshness of read data.
	//
	// For example, "replica" means eventually consistent reads, while
	// "master" supports transactions and gives you read-after-write consistency.
	//
	// Default: "master"
	// Allowed values: "master", "replica", "rdonly"
	TabletType string `json:"tablet_type"`

	// Streaming is true when streaming RPCs are used.
	// Recommended for large results.
	// Default: false
	Streaming bool

	// Timeout after which a pending query will be aborted.
	Timeout time.Duration
}

func newDefaultConfiguration() Configuration {
	return Configuration{
		Protocol:   "grpc",
		TabletType: "master",
		Streaming:  false,
	}
}

type conn struct {
	Configuration
	// tabletTypeProto is the protobof enum value of the string Configuration.TabletType.
	tabletTypeProto topodatapb.TabletType
	vtgateConn      *vtgateconn.VTGateConn
	tx              *vtgateconn.VTGateTx
}

func (c *conn) dial() error {
	var err error
	if c.Protocol == "" {
		c.vtgateConn, err = vtgateconn.Dial(context.Background(), c.Address, c.Timeout)
	} else {
		c.vtgateConn, err = vtgateconn.DialProtocol(context.Background(), c.Protocol, c.Address, c.Timeout)
	}
	return err
}

func (c *conn) useExecuteShards() bool {
	return c.Keyspace != "" && c.Shard != ""
}

func (c *conn) Prepare(query string) (driver.Stmt, error) {
	return &stmt{c: c, query: query}, nil
}

func (c *conn) Close() error {
	c.vtgateConn.Close()
	return nil
}

func (c *conn) Begin() (driver.Tx, error) {
	if c.Streaming {
		return nil, errors.New("transaction not allowed for streaming connection")
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.Timeout)
	defer cancel()
	tx, err := c.vtgateConn.Begin(ctx)
	if err != nil {
		return nil, err
	}
	c.tx = tx
	return c, nil
}

func (c *conn) Commit() error {
	if c.tx == nil {
		return errors.New("commit: not in transaction")
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.Timeout)
	defer func() {
		cancel()
		c.tx = nil
	}()
	return c.tx.Commit(ctx)
}

func (c *conn) Rollback() error {
	if c.tx == nil {
		return nil
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.Timeout)
	defer func() {
		cancel()
		c.tx = nil
	}()
	return c.tx.Rollback(ctx)
}

type stmt struct {
	c     *conn
	query string
}

func (s *stmt) Close() error {
	return nil
}

func (s *stmt) NumInput() int {
	return -1
}

func (s *stmt) Exec(args []driver.Value) (driver.Result, error) {
	ctx, cancel := context.WithTimeout(context.Background(), s.c.Timeout)
	defer cancel()

	if s.c.Streaming {
		return nil, errors.New("Exec not allowed for streaming connections")
	}

	qr, err := s.executeVitess(ctx, args)
	if err != nil {
		return nil, err
	}
	return result{int64(qr.InsertID), int64(qr.RowsAffected)}, nil
}

func (s *stmt) Query(args []driver.Value) (driver.Rows, error) {
	ctx, cancel := context.WithTimeout(context.Background(), s.c.Timeout)

	if s.c.Streaming {
		var qrc <-chan *sqltypes.Result
		var errFunc vtgateconn.ErrFunc
		var err error
		if s.c.useExecuteShards() {
			qrc, errFunc, err = s.c.vtgateConn.StreamExecuteShards(ctx, s.query, s.c.Keyspace, []string{s.c.Shard}, makeBindVars(args), s.c.tabletTypeProto)
		} else {
			qrc, errFunc, err = s.c.vtgateConn.StreamExecute(ctx, s.query, makeBindVars(args), s.c.tabletTypeProto)
		}
		if err != nil {
			return nil, err
		}
		return newStreamingRows(qrc, errFunc, cancel), nil
	}
	// Do not cancel in case of a streaming query. It will do it itself.
	defer cancel()

	qr, err := s.executeVitess(ctx, args)
	if err != nil {
		return nil, err
	}
	return newRows(qr), nil
}

func (s *stmt) executeVitess(ctx context.Context, args []driver.Value) (*sqltypes.Result, error) {
	if s.c.tx != nil {
		if s.c.useExecuteShards() {
			return s.c.tx.ExecuteShards(ctx, s.query, s.c.Keyspace, []string{s.c.Shard}, makeBindVars(args), s.c.tabletTypeProto, false /* notInTransaction */)
		}
		return s.c.tx.Execute(ctx, s.query, makeBindVars(args), s.c.tabletTypeProto, false /* notInTransaction */)
	}

	// Non-transactional case.
	if s.c.useExecuteShards() {
		return s.c.vtgateConn.ExecuteShards(ctx, s.query, s.c.Keyspace, []string{s.c.Shard}, makeBindVars(args), s.c.tabletTypeProto)
	}
	return s.c.vtgateConn.Execute(ctx, s.query, makeBindVars(args), s.c.tabletTypeProto)
}

func makeBindVars(args []driver.Value) map[string]interface{} {
	if len(args) == 0 {
		return map[string]interface{}{}
	}
	bv := make(map[string]interface{}, len(args))
	for i, v := range args {
		bv[fmt.Sprintf("v%d", i+1)] = v
	}
	return bv
}

type result struct {
	insertid, rowsaffected int64
}

func (r result) LastInsertId() (int64, error) {
	return r.insertid, nil
}

func (r result) RowsAffected() (int64, error) {
	return r.rowsaffected, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package sqlite3_test

import (
	"github.com/conformal/btcdb"
	"github.com/conformal/btcdb/sqlite3"
	"os"
	"path/filepath"
	"testing"
)

func TestFailOperational(t *testing.T) {
	sqlite3.SetTestingT(t)
	failtestOperationalMode(t, dbTmDefault)
	failtestOperationalMode(t, dbTmNormal)
	failtestOperationalMode(t, dbTmFast)
	failtestOperationalMode(t, dbTmNoVerify)
}

func failtestOperationalMode(t *testing.T, mode int) {
	// simplified basic operation is:
	// 1) fetch block from remote server
	// 2) look up all txin (except coinbase in db)
	// 3) insert block

	// Ignore db remove errors since it means we didn't have an old one.
	dbname := "tstdbop1"
	_ = os.Remove(dbname)
	db, err := btcdb.CreateDB("sqlite", dbname)
	if err != nil {
		t.Errorf("Failed to open test database %v", err)
		return
	}
	defer os.Remove(dbname)
	defer db.Close()

	switch mode {
	case dbTmDefault: // default
		// no setup
	case dbTmNormal: // explicit normal
		db.SetDBInsertMode(btcdb.InsertNormal)
	case dbTmFast: // fast mode
		db.SetDBInsertMode(btcdb.InsertFast)
		if sqldb, ok := db.(*sqlite3.SqliteDb); ok {
			sqldb.TempTblMax = 100
		} else {
			t.Errorf("not right type")
		}
	case dbTmNoVerify: // validated block
		// no point in testing this
		return
	}

	// Since we are dealing with small dataset, reduce cache size
	sqlite3.SetBlockCacheSize(db, 2)
	sqlite3.SetTxCacheSize(db, 3)

	testdatafile := filepath.Join("testdata", "blocks1-256.bz2")
	blocks, err := loadBlocks(t, testdatafile)
	if err != nil {
		t.Errorf("Unable to load blocks from test data for mode %v: %v",
			mode, err)
		return
	}

	err = nil
out:
	for height := int64(0); height < int64(len(blocks)); height++ {
		block := blocks[height]

		mblock := block.MsgBlock()
		blockname, _ :=  block.Sha()

		if height == 248 {
			// time to corrupt the datbase, to see if it leaves the block or tx in the db
			if len(mblock.Transactions) != 2 {
				t.Errorf("transaction #248 should have two transactions txid %v ?= 828ef3b079f9c23829c56fe86e85b4a69d9e06e5b54ea597eef5fb3ffef509fe", blockname)
				return
			}
			tx := mblock.Transactions[1]
			txin := tx.TxIn[0]
			origintxsha := &txin.PreviousOutpoint.Hash
			sqlite3.KillTx(db, origintxsha)
			_, _, _, _, err = db.FetchTxAllBySha(origintxsha)
			if err == nil {
				t.Errorf("deleted tx found %v", origintxsha)
			}
		}

		if height == 248 {
		}
		newheight, err := db.InsertBlock(block)
		if err != nil {
			if height != 248 {
				t.Errorf("failed to insert block %v err %v", height, err)
				break out
			}
		} else {
			if height == 248 {
				t.Errorf("block insert with missing input tx succeeded block %v err %v", height, err)
				break out
			}
		}
		if height == 248 {
			for _, tx := range mblock.Transactions {
				txsha, err := tx.TxSha(block.ProtocolVersion())
				_, _, _, _, err = db.FetchTxAllBySha(&txsha)
				if err == nil {
					t.Errorf("referenced tx found, should not have been %v, ", txsha)
				}
			}
		}
		if height == 248 {
			exists := db.ExistsSha(blockname)
			if exists == true {
				t.Errorf("block still present after failed insert")
			}
			// if we got here with no error, testing was successful
			break out
		}
		if newheight != height {
			t.Errorf("height mismatch expect %v returned %v", height, newheight)
			break out
		}
	}

	switch mode {
	case dbTmDefault: // default
		// no cleanup
	case dbTmNormal: // explicit normal
		// no cleanup
	case dbTmFast: // fast mode
		db.SetDBInsertMode(btcdb.InsertNormal)
	case dbTmNoVerify: // validated block
		db.SetDBInsertMode(btcdb.InsertNormal)
	}
}
<commit_msg>Update the insertfail test for the latest changes.<commit_after>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package sqlite3_test

import (
	"github.com/conformal/btcdb"
	"github.com/conformal/btcdb/sqlite3"
	"os"
	"path/filepath"
	"testing"
)

func TestFailOperational(t *testing.T) {
	sqlite3.SetTestingT(t)
	failtestOperationalMode(t, dbTmDefault)
	failtestOperationalMode(t, dbTmNormal)
	failtestOperationalMode(t, dbTmFast)
	failtestOperationalMode(t, dbTmNoVerify)
}

func failtestOperationalMode(t *testing.T, mode int) {
	// simplified basic operation is:
	// 1) fetch block from remote server
	// 2) look up all txin (except coinbase in db)
	// 3) insert block

	// Ignore db remove errors since it means we didn't have an old one.
	dbname := "tstdbop1"
	_ = os.Remove(dbname)
	db, err := btcdb.CreateDB("sqlite", dbname)
	if err != nil {
		t.Errorf("Failed to open test database %v", err)
		return
	}
	defer os.Remove(dbname)
	defer db.Close()

	switch mode {
	case dbTmDefault: // default
		// no setup
	case dbTmNormal: // explicit normal
		db.SetDBInsertMode(btcdb.InsertNormal)
	case dbTmFast: // fast mode
		db.SetDBInsertMode(btcdb.InsertFast)
		if sqldb, ok := db.(*sqlite3.SqliteDb); ok {
			sqldb.TempTblMax = 100
		} else {
			t.Errorf("not right type")
		}
	case dbTmNoVerify: // validated block
		// no point in testing this
		return
	}

	// Since we are dealing with small dataset, reduce cache size
	sqlite3.SetBlockCacheSize(db, 2)
	sqlite3.SetTxCacheSize(db, 3)

	testdatafile := filepath.Join("testdata", "blocks1-256.bz2")
	blocks, err := loadBlocks(t, testdatafile)
	if err != nil {
		t.Errorf("Unable to load blocks from test data for mode %v: %v",
			mode, err)
		return
	}

	err = nil
out:
	for height := int64(0); height < int64(len(blocks)); height++ {
		block := blocks[height]

		mblock := block.MsgBlock()
		blockname, _ :=  block.Sha()

		if height == 248 {
			// time to corrupt the datbase, to see if it leaves the block or tx in the db
			if len(mblock.Transactions) != 2 {
				t.Errorf("transaction #248 should have two transactions txid %v ?= 828ef3b079f9c23829c56fe86e85b4a69d9e06e5b54ea597eef5fb3ffef509fe", blockname)
				return
			}
			tx := mblock.Transactions[1]
			txin := tx.TxIn[0]
			origintxsha := &txin.PreviousOutpoint.Hash
			sqlite3.KillTx(db, origintxsha)
			_, _, _, _, err = db.FetchTxAllBySha(origintxsha)
			if err == nil {
				t.Errorf("deleted tx found %v", origintxsha)
			}
		}

		if height == 248 {
		}
		newheight, err := db.InsertBlock(block)
		if err != nil {
			if height != 248 {
				t.Errorf("failed to insert block %v err %v", height, err)
				break out
			}
		} else {
			if height == 248 {
				t.Errorf("block insert with missing input tx succeeded block %v err %v", height, err)
				break out
			}
		}
		if height == 248 {
			for _, tx := range mblock.Transactions {
				txsha, err := tx.TxSha()
				_, _, _, _, err = db.FetchTxAllBySha(&txsha)
				if err == nil {
					t.Errorf("referenced tx found, should not have been %v, ", txsha)
				}
			}
		}
		if height == 248 {
			exists := db.ExistsSha(blockname)
			if exists == true {
				t.Errorf("block still present after failed insert")
			}
			// if we got here with no error, testing was successful
			break out
		}
		if newheight != height {
			t.Errorf("height mismatch expect %v returned %v", height, newheight)
			break out
		}
	}

	switch mode {
	case dbTmDefault: // default
		// no cleanup
	case dbTmNormal: // explicit normal
		// no cleanup
	case dbTmFast: // fast mode
		db.SetDBInsertMode(btcdb.InsertNormal)
	case dbTmNoVerify: // validated block
		db.SetDBInsertMode(btcdb.InsertNormal)
	}
}
<|endoftext|>


<commit_before>// Contains the Whisper protocol Envelope element. For formal details please see
// the specs at https://github.com/ethereum/wiki/wiki/Whisper-PoC-1-Protocol-Spec#envelopes.

package whisper

import (
	"crypto/ecdsa"
	"encoding/binary"
	"fmt"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/crypto/ecies"
	"github.com/ethereum/go-ethereum/rlp"
)

// Envelope represents a clear-text data packet to transmit through the Whisper
// network. Its contents may or may not be encrypted and signed.
type Envelope struct {
	Expiry uint32 // Whisper protocol specifies int32, really should be int64
	TTL    uint32 // ^^^^^^
	Topics []Topic
	Data   []byte
	Nonce  uint32

	hash common.Hash // Cached hash of the envelope to avoid rehashing every time
}

// NewEnvelope wraps a Whisper message with expiration and destination data
// included into an envelope for network forwarding.
func NewEnvelope(ttl time.Duration, topics []Topic, msg *Message) *Envelope {
	return &Envelope{
		Expiry: uint32(time.Now().Add(ttl).Unix()),
		TTL:    uint32(ttl.Seconds()),
		Topics: topics,
		Data:   msg.bytes(),
		Nonce:  0,
	}
}

// Seal closes the envelope by spending the requested amount of time as a proof
// of work on hashing the data.
func (self *Envelope) Seal(pow time.Duration) {
	d := make([]byte, 64)
	copy(d[:32], self.rlpWithoutNonce())

	finish, bestBit := time.Now().Add(pow).UnixNano(), 0
	for nonce := uint32(0); time.Now().UnixNano() < finish; {
		for i := 0; i < 1024; i++ {
			binary.BigEndian.PutUint32(d[60:], nonce)

			firstBit := common.FirstBitSet(common.BigD(crypto.Sha3(d)))
			if firstBit > bestBit {
				self.Nonce, bestBit = nonce, firstBit
			}
			nonce++
		}
	}
}

// rlpWithoutNonce returns the RLP encoded envelope contents, except the nonce.
func (self *Envelope) rlpWithoutNonce() []byte {
	enc, _ := rlp.EncodeToBytes([]interface{}{self.Expiry, self.TTL, self.Topics, self.Data})
	return enc
}

// Open extracts the message contained within a potentially encrypted envelope.
func (self *Envelope) Open(key *ecdsa.PrivateKey) (msg *Message, err error) {
	// Split open the payload into a message construct
	data := self.Data

	message := &Message{
		Flags: data[0],
	}
	data = data[1:]

	if message.Flags&signatureFlag == signatureFlag {
		if len(data) < signatureLength {
			return nil, fmt.Errorf("unable to open envelope. First bit set but len(data) < len(signature)")
		}
		message.Signature, data = data[:signatureLength], data[signatureLength:]
	}
	message.Payload = data

	// Decrypt the message, if requested
	if key == nil {
		return message, nil
	}
	err = message.decrypt(key)
	switch err {
	case nil:
		return message, nil

	case ecies.ErrInvalidPublicKey: // Payload isn't encrypted
		return message, err

	default:
		return nil, fmt.Errorf("unable to open envelope, decrypt failed: %v", err)
	}
}

// Hash returns the SHA3 hash of the envelope, calculating it if not yet done.
func (self *Envelope) Hash() common.Hash {
	if (self.hash == common.Hash{}) {
		enc, _ := rlp.EncodeToBytes(self)
		self.hash = crypto.Sha3Hash(enc)
	}
	return self.hash
}

// rlpenv is an Envelope but is not an rlp.Decoder.
// It is used for decoding because we need to
type rlpenv Envelope

// DecodeRLP decodes an Envelope from an RLP data stream.
func (self *Envelope) DecodeRLP(s *rlp.Stream) error {
	raw, err := s.Raw()
	if err != nil {
		return err
	}
	if err := rlp.DecodeBytes(raw, (*rlpenv)(self)); err != nil {
		return err
	}
	self.hash = crypto.Sha3Hash(raw)
	return nil
}
<commit_msg>whisper: fix comment for rlpenv<commit_after>// Contains the Whisper protocol Envelope element. For formal details please see
// the specs at https://github.com/ethereum/wiki/wiki/Whisper-PoC-1-Protocol-Spec#envelopes.

package whisper

import (
	"crypto/ecdsa"
	"encoding/binary"
	"fmt"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/crypto/ecies"
	"github.com/ethereum/go-ethereum/rlp"
)

// Envelope represents a clear-text data packet to transmit through the Whisper
// network. Its contents may or may not be encrypted and signed.
type Envelope struct {
	Expiry uint32 // Whisper protocol specifies int32, really should be int64
	TTL    uint32 // ^^^^^^
	Topics []Topic
	Data   []byte
	Nonce  uint32

	hash common.Hash // Cached hash of the envelope to avoid rehashing every time
}

// NewEnvelope wraps a Whisper message with expiration and destination data
// included into an envelope for network forwarding.
func NewEnvelope(ttl time.Duration, topics []Topic, msg *Message) *Envelope {
	return &Envelope{
		Expiry: uint32(time.Now().Add(ttl).Unix()),
		TTL:    uint32(ttl.Seconds()),
		Topics: topics,
		Data:   msg.bytes(),
		Nonce:  0,
	}
}

// Seal closes the envelope by spending the requested amount of time as a proof
// of work on hashing the data.
func (self *Envelope) Seal(pow time.Duration) {
	d := make([]byte, 64)
	copy(d[:32], self.rlpWithoutNonce())

	finish, bestBit := time.Now().Add(pow).UnixNano(), 0
	for nonce := uint32(0); time.Now().UnixNano() < finish; {
		for i := 0; i < 1024; i++ {
			binary.BigEndian.PutUint32(d[60:], nonce)

			firstBit := common.FirstBitSet(common.BigD(crypto.Sha3(d)))
			if firstBit > bestBit {
				self.Nonce, bestBit = nonce, firstBit
			}
			nonce++
		}
	}
}

// rlpWithoutNonce returns the RLP encoded envelope contents, except the nonce.
func (self *Envelope) rlpWithoutNonce() []byte {
	enc, _ := rlp.EncodeToBytes([]interface{}{self.Expiry, self.TTL, self.Topics, self.Data})
	return enc
}

// Open extracts the message contained within a potentially encrypted envelope.
func (self *Envelope) Open(key *ecdsa.PrivateKey) (msg *Message, err error) {
	// Split open the payload into a message construct
	data := self.Data

	message := &Message{
		Flags: data[0],
	}
	data = data[1:]

	if message.Flags&signatureFlag == signatureFlag {
		if len(data) < signatureLength {
			return nil, fmt.Errorf("unable to open envelope. First bit set but len(data) < len(signature)")
		}
		message.Signature, data = data[:signatureLength], data[signatureLength:]
	}
	message.Payload = data

	// Decrypt the message, if requested
	if key == nil {
		return message, nil
	}
	err = message.decrypt(key)
	switch err {
	case nil:
		return message, nil

	case ecies.ErrInvalidPublicKey: // Payload isn't encrypted
		return message, err

	default:
		return nil, fmt.Errorf("unable to open envelope, decrypt failed: %v", err)
	}
}

// Hash returns the SHA3 hash of the envelope, calculating it if not yet done.
func (self *Envelope) Hash() common.Hash {
	if (self.hash == common.Hash{}) {
		enc, _ := rlp.EncodeToBytes(self)
		self.hash = crypto.Sha3Hash(enc)
	}
	return self.hash
}

// DecodeRLP decodes an Envelope from an RLP data stream.
func (self *Envelope) DecodeRLP(s *rlp.Stream) error {
	raw, err := s.Raw()
	if err != nil {
		return err
	}
	// The decoding of Envelope uses the struct fields but also needs
	// to compute the hash of the whole RLP-encoded envelope. This
	// type has the same structure as Envelope but is not an
	// rlp.Decoder so we can reuse the Envelope struct definition.
	type rlpenv Envelope
	if err := rlp.DecodeBytes(raw, (*rlpenv)(self)); err != nil {
		return err
	}
	self.hash = crypto.Sha3Hash(raw)
	return nil
}
<|endoftext|>