<commit_before>package google

import (
	"fmt"
	"testing"

	"github.com/hashicorp/terraform-plugin-sdk/helper/acctest"
	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/terraform"
	"google.golang.org/api/bigquery/v2"
)

func TestAccBigQueryDataset_basic(t *testing.T) {
	t.Parallel()

	datasetID := fmt.Sprintf("tf_test_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckBigQueryDatasetDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryDataset(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
			{
				Config: testAccBigQueryDatasetUpdated(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func TestAccBigQueryDataset_datasetWithContents(t *testing.T) {
	t.Parallel()

	datasetID := fmt.Sprintf("tf_test_%s", acctest.RandString(10))
	tableID := fmt.Sprintf("tf_test_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckBigQueryDatasetDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryDatasetDeleteContents(datasetID),
				Check:  testAccAddTable(datasetID, tableID),
			},
			{
				ResourceName:            "google_bigquery_dataset.contents_test",
				ImportState:             true,
				ImportStateVerify:       true,
				ImportStateVerifyIgnore: []string{"delete_contents_on_destroy"},
			},
		},
	})
}

func TestAccBigQueryDataset_access(t *testing.T) {
	t.Parallel()

	datasetID := fmt.Sprintf("tf_test_access_%s", acctest.RandString(10))
	otherDatasetID := fmt.Sprintf("tf_test_other_%s", acctest.RandString(10))
	otherTableID := fmt.Sprintf("tf_test_other_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckBigQueryDatasetDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryDatasetWithOneAccess(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.access_test",
				ImportState:       true,
				ImportStateVerify: true,
			},
			{
				Config: testAccBigQueryDatasetWithTwoAccess(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.access_test",
				ImportState:       true,
				ImportStateVerify: true,
			},
			{
				Config: testAccBigQueryDatasetWithOneAccess(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.access_test",
				ImportState:       true,
				ImportStateVerify: true,
			},
			{
				Config: testAccBigQueryDatasetWithViewAccess(datasetID, otherDatasetID, otherTableID),
			},
			{
				ResourceName:      "google_bigquery_dataset.access_test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func TestAccBigQueryDataset_regionalLocation(t *testing.T) {
	t.Parallel()

	datasetID1 := fmt.Sprintf("tf_test_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckBigQueryDatasetDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryRegionalDataset(datasetID1, "asia-south1"),
			},
			{
				ResourceName:      "google_bigquery_dataset.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func TestAccBigQueryDataset_cmek(t *testing.T) {
	t.Parallel()

	kms := BootstrapKMSKeyInLocation(t, "us")
	pid := getTestProjectFromEnv()
	datasetID1 := fmt.Sprintf("tf_test_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryDataset_cmek(pid, datasetID1, kms.CryptoKey.Name),
			},
			{
				ResourceName:      "google_bigquery_dataset.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func testAccAddTable(datasetID string, tableID string) resource.TestCheckFunc {
	// Not actually a check, but adds a table independently of terraform
	return func(s *terraform.State) error {
		config := testAccProvider.Meta().(*Config)
		table := &bigquery.Table{
			TableReference: &bigquery.TableReference{
				DatasetId: datasetID,
				TableId:   tableID,
				ProjectId: config.Project,
			},
		}
		_, err := config.clientBigQuery.Tables.Insert(config.Project, datasetID, table).Do()
		if err != nil {
			return fmt.Errorf("Could not create table")
		}
		return nil
	}
}

func testAccBigQueryDataset(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "test" {
  dataset_id                      = "%s"
  friendly_name                   = "foo"
  description                     = "This is a foo description"
  location                        = "EU"
  default_partition_expiration_ms = 3600000
  default_table_expiration_ms     = 3600000

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID)
}

func testAccBigQueryDatasetUpdated(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "test" {
  dataset_id                      = "%s"
  friendly_name                   = "bar"
  description                     = "This is a bar description"
  location                        = "EU"
  default_partition_expiration_ms = 7200000
  default_table_expiration_ms     = 7200000

  labels = {
    env                         = "bar"
    default_table_expiration_ms = 7200000
  }
}
`, datasetID)
}

func testAccBigQueryDatasetDeleteContents(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "contents_test" {
  dataset_id                      = "%s"
  friendly_name                   = "foo"
  description                     = "This is a foo description"
  location                        = "EU"
  default_partition_expiration_ms = 3600000
  default_table_expiration_ms     = 3600000
  delete_contents_on_destroy      = true

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID)
}

func testAccBigQueryRegionalDataset(datasetID string, location string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "test" {
  dataset_id                  = "%s"
  friendly_name               = "foo"
  description                 = "This is a foo description"
  location                    = "%s"
  default_table_expiration_ms = 3600000

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID, location)
}

func testAccBigQueryDatasetWithOneAccess(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "access_test" {
  dataset_id = "%s"

  access {
    role          = "OWNER"
    user_by_email = "Joe@example.com"
  }

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID)
}

func testAccBigQueryDatasetWithTwoAccess(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "access_test" {
  dataset_id = "%s"

  access {
    role          = "OWNER"
    user_by_email = "Joe@example.com"
  }
  access {
    role   = "READER"
    domain = "example.com"
  }

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID)
}

func testAccBigQueryDatasetWithViewAccess(datasetID, otherDatasetID, otherTableID string) string {
	// Note that we have to add a non-view access to prevent BQ from creating 4 default
	// access entries.
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "other_dataset" {
  dataset_id = "%s"
}

resource "google_bigquery_table" "table_with_view" {
  table_id   = "%s"
  dataset_id = google_bigquery_dataset.other_dataset.dataset_id

  time_partitioning {
    type = "DAY"
  }

  view {
    query          = "SELECT state FROM [lookerdata:cdc.project_tycho_reports]"
    use_legacy_sql = true
  }
}

resource "google_bigquery_dataset" "access_test" {
  dataset_id = "%s"

  access {
    role          = "OWNER"
    user_by_email = "Joe@example.com"
  }
  access {
    view {
      project_id = google_bigquery_dataset.other_dataset.project
      dataset_id = google_bigquery_dataset.other_dataset.dataset_id
      table_id   = google_bigquery_table.table_with_view.table_id
    }
  }

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, otherDatasetID, otherTableID, datasetID)
}

func testAccBigQueryDataset_cmek(pid, datasetID, kmsKey string) string {
	return fmt.Sprintf(`
data "google_project" "project" {
  project_id = "%s"
}

resource "google_project_iam_member" "kms-project-binding" {
  project = data.google_project.project.project_id
  role    = "roles/cloudkms.cryptoKeyEncrypterDecrypter"
  member  = "serviceAccount:bq-${data.google_project.project.number}@bigquery-encryption.iam.gserviceaccount.com"
}

resource "google_bigquery_dataset" "test" {
  dataset_id                  = "%s"
  friendly_name               = "test"
  description                 = "This is a test description"
  location                    = "US"
  default_table_expiration_ms = 3600000

  default_encryption_configuration {
    kms_key_name = "%s"
  }

  project = google_project_iam_member.kms-project-binding.project
}
`, pid, datasetID, kmsKey)
}
<commit_msg>Tests are failing due to domain verification (#3176)<commit_after>package google

import (
	"fmt"
	"testing"

	"github.com/hashicorp/terraform-plugin-sdk/helper/acctest"
	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/terraform"
	"google.golang.org/api/bigquery/v2"
)

func TestAccBigQueryDataset_basic(t *testing.T) {
	t.Parallel()

	datasetID := fmt.Sprintf("tf_test_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckBigQueryDatasetDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryDataset(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
			{
				Config: testAccBigQueryDatasetUpdated(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func TestAccBigQueryDataset_datasetWithContents(t *testing.T) {
	t.Parallel()

	datasetID := fmt.Sprintf("tf_test_%s", acctest.RandString(10))
	tableID := fmt.Sprintf("tf_test_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckBigQueryDatasetDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryDatasetDeleteContents(datasetID),
				Check:  testAccAddTable(datasetID, tableID),
			},
			{
				ResourceName:            "google_bigquery_dataset.contents_test",
				ImportState:             true,
				ImportStateVerify:       true,
				ImportStateVerifyIgnore: []string{"delete_contents_on_destroy"},
			},
		},
	})
}

func TestAccBigQueryDataset_access(t *testing.T) {
	t.Parallel()

	datasetID := fmt.Sprintf("tf_test_access_%s", acctest.RandString(10))
	otherDatasetID := fmt.Sprintf("tf_test_other_%s", acctest.RandString(10))
	otherTableID := fmt.Sprintf("tf_test_other_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckBigQueryDatasetDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryDatasetWithOneAccess(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.access_test",
				ImportState:       true,
				ImportStateVerify: true,
			},
			{
				Config: testAccBigQueryDatasetWithTwoAccess(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.access_test",
				ImportState:       true,
				ImportStateVerify: true,
			},
			{
				Config: testAccBigQueryDatasetWithOneAccess(datasetID),
			},
			{
				ResourceName:      "google_bigquery_dataset.access_test",
				ImportState:       true,
				ImportStateVerify: true,
			},
			{
				Config: testAccBigQueryDatasetWithViewAccess(datasetID, otherDatasetID, otherTableID),
			},
			{
				ResourceName:      "google_bigquery_dataset.access_test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func TestAccBigQueryDataset_regionalLocation(t *testing.T) {
	t.Parallel()

	datasetID1 := fmt.Sprintf("tf_test_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckBigQueryDatasetDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryRegionalDataset(datasetID1, "asia-south1"),
			},
			{
				ResourceName:      "google_bigquery_dataset.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func TestAccBigQueryDataset_cmek(t *testing.T) {
	t.Parallel()

	kms := BootstrapKMSKeyInLocation(t, "us")
	pid := getTestProjectFromEnv()
	datasetID1 := fmt.Sprintf("tf_test_%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccBigQueryDataset_cmek(pid, datasetID1, kms.CryptoKey.Name),
			},
			{
				ResourceName:      "google_bigquery_dataset.test",
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func testAccAddTable(datasetID string, tableID string) resource.TestCheckFunc {
	// Not actually a check, but adds a table independently of terraform
	return func(s *terraform.State) error {
		config := testAccProvider.Meta().(*Config)
		table := &bigquery.Table{
			TableReference: &bigquery.TableReference{
				DatasetId: datasetID,
				TableId:   tableID,
				ProjectId: config.Project,
			},
		}
		_, err := config.clientBigQuery.Tables.Insert(config.Project, datasetID, table).Do()
		if err != nil {
			return fmt.Errorf("Could not create table")
		}
		return nil
	}
}

func testAccBigQueryDataset(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "test" {
  dataset_id                      = "%s"
  friendly_name                   = "foo"
  description                     = "This is a foo description"
  location                        = "EU"
  default_partition_expiration_ms = 3600000
  default_table_expiration_ms     = 3600000

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID)
}

func testAccBigQueryDatasetUpdated(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "test" {
  dataset_id                      = "%s"
  friendly_name                   = "bar"
  description                     = "This is a bar description"
  location                        = "EU"
  default_partition_expiration_ms = 7200000
  default_table_expiration_ms     = 7200000

  labels = {
    env                         = "bar"
    default_table_expiration_ms = 7200000
  }
}
`, datasetID)
}

func testAccBigQueryDatasetDeleteContents(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "contents_test" {
  dataset_id                      = "%s"
  friendly_name                   = "foo"
  description                     = "This is a foo description"
  location                        = "EU"
  default_partition_expiration_ms = 3600000
  default_table_expiration_ms     = 3600000
  delete_contents_on_destroy      = true

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID)
}

func testAccBigQueryRegionalDataset(datasetID string, location string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "test" {
  dataset_id                  = "%s"
  friendly_name               = "foo"
  description                 = "This is a foo description"
  location                    = "%s"
  default_table_expiration_ms = 3600000

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID, location)
}

func testAccBigQueryDatasetWithOneAccess(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "access_test" {
  dataset_id = "%s"

  access {
    role          = "OWNER"
    user_by_email = "Joe@example.com"
  }

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID)
}

func testAccBigQueryDatasetWithTwoAccess(datasetID string) string {
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "access_test" {
  dataset_id = "%s"

  access {
    role          = "OWNER"
    user_by_email = "Joe@example.com"
  }
  access {
    role   = "READER"
    domain = "hashicorp.com"
  }

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, datasetID)
}

func testAccBigQueryDatasetWithViewAccess(datasetID, otherDatasetID, otherTableID string) string {
	// Note that we have to add a non-view access to prevent BQ from creating 4 default
	// access entries.
	return fmt.Sprintf(`
resource "google_bigquery_dataset" "other_dataset" {
  dataset_id = "%s"
}

resource "google_bigquery_table" "table_with_view" {
  table_id   = "%s"
  dataset_id = google_bigquery_dataset.other_dataset.dataset_id

  time_partitioning {
    type = "DAY"
  }

  view {
    query          = "SELECT state FROM [lookerdata:cdc.project_tycho_reports]"
    use_legacy_sql = true
  }
}

resource "google_bigquery_dataset" "access_test" {
  dataset_id = "%s"

  access {
    role          = "OWNER"
    user_by_email = "Joe@example.com"
  }
  access {
    view {
      project_id = google_bigquery_dataset.other_dataset.project
      dataset_id = google_bigquery_dataset.other_dataset.dataset_id
      table_id   = google_bigquery_table.table_with_view.table_id
    }
  }

  labels = {
    env                         = "foo"
    default_table_expiration_ms = 3600000
  }
}
`, otherDatasetID, otherTableID, datasetID)
}

func testAccBigQueryDataset_cmek(pid, datasetID, kmsKey string) string {
	return fmt.Sprintf(`
data "google_project" "project" {
  project_id = "%s"
}

resource "google_project_iam_member" "kms-project-binding" {
  project = data.google_project.project.project_id
  role    = "roles/cloudkms.cryptoKeyEncrypterDecrypter"
  member  = "serviceAccount:bq-${data.google_project.project.number}@bigquery-encryption.iam.gserviceaccount.com"
}

resource "google_bigquery_dataset" "test" {
  dataset_id                  = "%s"
  friendly_name               = "test"
  description                 = "This is a test description"
  location                    = "US"
  default_table_expiration_ms = 3600000

  default_encryption_configuration {
    kms_key_name = "%s"
  }

  project = google_project_iam_member.kms-project-binding.project
}
`, pid, datasetID, kmsKey)
}
<|endoftext|>


<commit_before>package geolookup

import (
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	geo "github.com/getlantern/geolookup"
	"github.com/getlantern/golog"

	"github.com/getlantern/flashlight/pubsub"
	"github.com/getlantern/flashlight/ui"
)

const (
	messageType = `GeoLookup`

	basePublishSeconds     = 30
	publishSecondsVariance = basePublishSeconds - 10
	retryWaitMillis        = 100
)

var (
	log = golog.LoggerFor("flashlight.geolookup")

	service  *ui.Service
	client   atomic.Value
	cfgMutex sync.Mutex
	country  = atomicString()
	ip       = atomicString()
)

func atomicString() atomic.Value {
	var val atomic.Value
	val.Store("")
	return val
}

func GetIp() string {
	return ip.Load().(string)
}

func GetCountry() string {
	return country.Load().(string)
}

// Configure configures geolookup to use the given http.Client to perform
// lookups. geolookup runs in a continuous loop, periodically updating its
// location and publishing updates to any connected clients. We do this
// continually in order to detect when the computer's location has changed.
func Configure(newClient *http.Client) {
	cfgMutex.Lock()
	defer cfgMutex.Unlock()

	// Avoid annoying checks for nil later.
	ip.Store("")
	country.Store("")

	client.Store(newClient)

	if service == nil {
		err := registerService()
		if err != nil {
			log.Errorf("Unable to register service: %s", err)
			return
		}
		go write()
		go read()
		log.Debug("Running")
	}
}

func registerService() error {
	helloFn := func(write func(interface{}) error) error {
		country := GetCountry()
		if country == "" {
			log.Trace("No lastKnownCountry, not sending anything to client")
			return nil
		}
		log.Trace("Sending last known location to new client")
		return write(country)
	}

	var err error
	service, err = ui.Register(messageType, nil, helloFn)
	return err
}

func lookupIp(httpClient *http.Client) (string, string, error) {
	city, ip, err := geo.LookupIPWithClient("", httpClient)

	if err != nil {
		return "", "", err
	}
	return city.Country.IsoCode, ip, nil

	/*
		httpClient.Timeout = 60 * time.Second

		var err error
		var req *http.Request
		var resp *http.Response

		// Note this will typically be an HTTP client that uses direct domain fronting to
		// hit our server pool in the Netherlands.
		if req, err = http.NewRequest("HEAD", "http://nl.fallbacks.getiantem.org", nil); err != nil {
			return "", "", fmt.Errorf("Could not create request: %q", err)
		}

		if resp, err = httpClient.Do(req); err != nil {
			return "", "", fmt.Errorf("Could not get response from server: %q", err)
		}
		defer func() {
			if err := resp.Body.Close(); err != nil {
				log.Debugf("Unable to close reponse body: %v", err)
			}
		}()

		if resp.StatusCode != http.StatusOK {
			if full, err := httputil.DumpResponse(resp, true); err != nil {
				log.Errorf("Could not read full response %v", err)
			} else {
				log.Errorf("Unexpected response to geo IP lookup: %v", string(full))
			}
			return "", "", fmt.Errorf("Unexpected response status %d", resp.StatusCode)
		}

		ip := resp.Header.Get("Lantern-Ip")
		country := resp.Header.Get("Lantern-Country")

		log.Debugf("Got IP and country: %v, %v", ip, country)
		return country, ip, nil
	*/
}

func write() {
	consecutiveFailures := 0

	for {
		// Wait a random amount of time (to avoid looking too suspicious)
		// Note - rand was seeded with the startup time in flashlight.go
		n := rand.Intn(publishSecondsVariance)
		wait := time.Duration(basePublishSeconds-publishSecondsVariance/2+n) * time.Second

		oldIp := GetIp()
		oldCountry := GetCountry()

		newCountry, newIp, err := lookupIp(client.Load().(*http.Client))
		if err == nil {
			consecutiveFailures = 0
			if newIp != oldIp {
				log.Debugf("IP changed")
				ip.Store(newIp)
			}
			// Always publish location, even if unchanged
			pubsub.Pub(pubsub.IP, newIp)
			service.Out <- newCountry
		} else {
			msg := fmt.Sprintf("Unable to get current location: %s", err)
			// When retrying after a failure, wait a different amount of time
			retryWait := time.Duration(math.Pow(2, float64(consecutiveFailures))*float64(retryWaitMillis)) * time.Millisecond
			if retryWait < wait {
				log.Debug(msg)
				wait = retryWait
			} else {
				log.Error(msg)
			}
			log.Debugf("Waiting %v before retrying", wait)
			consecutiveFailures += 1
			// If available, publish last known location
			if oldCountry != "" {
				service.Out <- oldCountry
			}
		}

		time.Sleep(wait)
	}
}

func read() {
	for _ = range service.In {
		// Discard message, just in case any message is sent to this service.
	}
}
<commit_msg>don't publish ip every time adn don't check it so frequently<commit_after>package geolookup

import (
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	geo "github.com/getlantern/geolookup"
	"github.com/getlantern/golog"

	"github.com/getlantern/flashlight/pubsub"
	"github.com/getlantern/flashlight/ui"
)

const (
	messageType = `GeoLookup`

	basePublishSeconds     = 120
	publishSecondsVariance = basePublishSeconds - 120
	retryWaitMillis        = 100
)

var (
	log = golog.LoggerFor("flashlight.geolookup")

	service  *ui.Service
	client   atomic.Value
	cfgMutex sync.Mutex
	country  = atomicString()
	ip       = atomicString()
)

func atomicString() atomic.Value {
	var val atomic.Value
	val.Store("")
	return val
}

func GetIp() string {
	return ip.Load().(string)
}

func GetCountry() string {
	return country.Load().(string)
}

// Configure configures geolookup to use the given http.Client to perform
// lookups. geolookup runs in a continuous loop, periodically updating its
// location and publishing updates to any connected clients. We do this
// continually in order to detect when the computer's location has changed.
func Configure(newClient *http.Client) {
	cfgMutex.Lock()
	defer cfgMutex.Unlock()

	// Avoid annoying checks for nil later.
	ip.Store("")
	country.Store("")

	client.Store(newClient)

	if service == nil {
		err := registerService()
		if err != nil {
			log.Errorf("Unable to register service: %s", err)
			return
		}
		go write()
		go read()
		log.Debug("Running")
	}
}

func registerService() error {
	helloFn := func(write func(interface{}) error) error {
		country := GetCountry()
		if country == "" {
			log.Trace("No lastKnownCountry, not sending anything to client")
			return nil
		}
		log.Trace("Sending last known location to new client")
		return write(country)
	}

	var err error
	service, err = ui.Register(messageType, nil, helloFn)
	return err
}

func lookupIp(httpClient *http.Client) (string, string, error) {
	city, ip, err := geo.LookupIPWithClient("", httpClient)

	if err != nil {
		log.Errorf("Could not lookup IP %v", err)
		return "", "", err
	}
	return city.Country.IsoCode, ip, nil
}

func write() {
	consecutiveFailures := 0

	for {
		// Wait a random amount of time (to avoid looking too suspicious)
		// Note - rand was seeded with the startup time in flashlight.go
		n := rand.Intn(publishSecondsVariance)
		wait := time.Duration(basePublishSeconds-publishSecondsVariance/2+n) * time.Second

		oldIp := GetIp()
		oldCountry := GetCountry()

		newCountry, newIp, err := lookupIp(client.Load().(*http.Client))
		if err == nil {
			consecutiveFailures = 0
			if newIp != oldIp {
				log.Debugf("IP changed")
				ip.Store(newIp)
				pubsub.Pub(pubsub.IP, newIp)
			}
			// Always publish location, even if unchanged
			service.Out <- newCountry
		} else {
			msg := fmt.Sprintf("Unable to get current location: %s", err)
			// When retrying after a failure, wait a different amount of time
			retryWait := time.Duration(math.Pow(2, float64(consecutiveFailures))*float64(retryWaitMillis)) * time.Millisecond
			if retryWait < wait {
				log.Debug(msg)
				wait = retryWait
			} else {
				log.Error(msg)
			}
			log.Debugf("Waiting %v before retrying", wait)
			consecutiveFailures += 1
			// If available, publish last known location
			if oldCountry != "" {
				service.Out <- oldCountry
			}
		}

		time.Sleep(wait)
	}
}

func read() {
	for _ = range service.In {
		// Discard message, just in case any message is sent to this service.
	}
}
<|endoftext|>


<commit_before>package retry

import (
	"context"
	"fmt"
	"testing"
	"time"
)

func TestRetrySuccess(t *testing.T) {
	cnt := 0
	err := Retry(3, 1*time.Nanosecond, func() error {
		if cnt == 0 {
			cnt++
			return fmt.Errorf("retry")
		}
		return nil
	})

	if err != nil {
		t.Errorf("error should be nil, %s", err)
	}

	if cnt != 1 {
		t.Errorf("cnt should be 1, %d", cnt)
	}
}

func TestRetryFail(t *testing.T) {
	cnt := 0
	err := Retry(4, 1*time.Nanosecond, func() error {
		cnt++
		return fmt.Errorf("retry")
	})

	if err == nil || err.Error() != "retry" {
		t.Errorf("error should be occured")
	}

	if cnt != 4 {
		t.Errorf("cnt should be 4, but %d", cnt)
	}
}

func TestRetryWithContext(t *testing.T) {
	cnt := 0
	ctx, _ := context.WithTimeout(context.Background(), 150*time.Millisecond)
	err := RetryWithContext(ctx, 3, 100*time.Millisecond, func() error {
		cnt++
		return fmt.Errorf("retry")
	})

	if err != context.DeadlineExceeded {
		t.Errorf("error should be %s, %s", context.DeadlineExceeded, err)
	}

	if cnt != 2 {
		t.Errorf("cnt should be 2, %d", cnt)
	}
}
<commit_msg>call cancel of the context in the test<commit_after>package retry

import (
	"context"
	"fmt"
	"testing"
	"time"
)

func TestRetrySuccess(t *testing.T) {
	cnt := 0
	err := Retry(3, 1*time.Nanosecond, func() error {
		if cnt == 0 {
			cnt++
			return fmt.Errorf("retry")
		}
		return nil
	})

	if err != nil {
		t.Errorf("error should be nil, %s", err)
	}

	if cnt != 1 {
		t.Errorf("cnt should be 1, %d", cnt)
	}
}

func TestRetryFail(t *testing.T) {
	cnt := 0
	err := Retry(4, 1*time.Nanosecond, func() error {
		cnt++
		return fmt.Errorf("retry")
	})

	if err == nil || err.Error() != "retry" {
		t.Errorf("error should be occured")
	}

	if cnt != 4 {
		t.Errorf("cnt should be 4, but %d", cnt)
	}
}

func TestRetryWithContext(t *testing.T) {
	cnt := 0
	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
	defer cancel()
	err := RetryWithContext(ctx, 3, 100*time.Millisecond, func() error {
		cnt++
		return fmt.Errorf("retry")
	})

	if err != context.DeadlineExceeded {
		t.Errorf("error should be %s, %s", context.DeadlineExceeded, err)
	}

	if cnt != 2 {
		t.Errorf("cnt should be 2, %d", cnt)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"database/sql"
	"log"
	"net/http"
	"sync"
	"time"
)

type RoomNameMap map[RoomID]string
type RoomGroupMap map[BuildingName]map[FloorID][]RoomID

const (
	INTERVAL = 1 * time.Minute
)

type RoomStatus struct {
	RoomID  RoomID         `json:"id"`
	Sensors []SensorStatus `json:"sensors"`

	Hot     uint64 `json:"hot"`
	Comfort uint64 `json:"comfort"`
	Cold    uint64 `json:"cold"`
	lock    sync.RWMutex
}

type MyVote struct {
	Vote      VoteChoice `json:"vote"`
	Timestamp int64      `json:"timestamp"`
}

type RoomStatusManager struct {
	db        *sql.DB
	thingworx *ThingWorxClient

	sensorCache map[RoomID]map[ThingName]SensorStatus
	cacheLock   sync.RWMutex
}

type RoomStatusTx struct {
	rsm *RoomStatusManager
	tx  *sql.Tx
	s   *Session
}

type SensorStatus struct {
	Temperature float64 `json:"temperature"`
	Humidity    float64 `json:"humidity"`
	IsConnected bool    `json:"isConnected"`

	expire time.Time
}

func NewRoomStatusManager(db *sql.DB, thingworx *ThingWorxClient, ctx context.Context) *RoomStatusManager {
	// create RSM
	rs := &RoomStatusManager{}
	rs.db = db
	rs.thingworx = thingworx
	rs.sensorCache = make(map[RoomID]map[ThingName]SensorStatus)

	go rs.cacheUpdater(ctx)
	return rs
}

func (rsm *RoomStatusManager) GetTx(w http.ResponseWriter, req *http.Request) (*RoomStatusTx, error) {
	tx, err := rsm.db.Begin()
	if err != nil {
		return nil, err
	}
	s := GetSession(w, req, tx)
	return &RoomStatusTx{
		rsm: rsm,
		tx:  tx,
		s:   s,
	}, nil
}

func (rst *RoomStatusTx) Rollback() error {
	return rst.tx.Rollback()
}

func (rst *RoomStatusTx) Commit() error {
	return rst.tx.Commit()
}

func (rst *RoomStatusTx) GetRoomName(id RoomID) (name string, err error) {
	err = rst.tx.QueryRow(
		`SELECT name FROM room
		WHERE room_id=?`,
		id,
	).Scan(&name)
	return
}

// 投票内容を取得する。未投票の場合はnilを返す
func (rst *RoomStatusTx) GetMyVote(id RoomID) (vote *MyVote, err error) {
	var v Vote
	if err = rst.tx.QueryRow(
		`SELECT choice, timestamp FROM vote
			WHERE session_id=? AND room_id=?`,
		rst.s.SessionID, id,
	).Scan(&v.Choice, &v.Timestamp); err != nil {
		if err == sql.ErrNoRows {
			// 未投票の状態。
			err = nil
			return
		}
		return
	}

	vote = &MyVote{
		Vote:      v.Choice,
		Timestamp: v.Timestamp.Unix(),
	}
	return vote, err
}

func (rst *RoomStatusTx) GetStatus(id RoomID) (*RoomStatus, error) {
	rs := &RoomStatus{
		RoomID: id,
	}

	var ok bool
	rs.Sensors, ok = rst.rsm.getSensorStatusFromCache(id)
	if !ok {
		// センサーの状態を更新できていない状態。
		rs.Sensors = []SensorStatus{}
	}

	rows, err := rst.tx.Query(
		`SELECT choice, count(vote_id) FROM vote
		WHERE room_id=?
		GROUP BY choice`,
		id,
	)
	if err != nil {
		return nil, err
	}
	for rows.Next() {
		var choice VoteChoice
		var count uint64
		if err := rows.Scan(&choice, &count); err != nil {
			return nil, err
		}
		switch choice {
		case Hot:
			rs.Hot = count
		case Comfort:
			rs.Comfort = count
		case Cold:
			rs.Cold = count
		}
	}
	return rs, nil
}

func (rst *RoomStatusTx) Vote(id RoomID, choice VoteChoice) error {
	vote := Vote{
		RoomID: id,
		S:      rst.s,
	}

	if err := rst.tx.QueryRow(
		`SELECT vote_id FROM vote
		WHERE session_id=? AND room_id=?`,
		rst.s.SessionID, id,
	).Scan(&vote.VoteID); err != nil {
		switch err {
		case sql.ErrNoRows:
			// 未投票であることを表す、0を代入
			vote.VoteID = VoteID(0)
		default:
			return err
		}
	}

	return vote.UpdateChoice(rst.tx, choice)
}

func (rst *RoomStatusTx) GetAllRoomsInfo() (names RoomNameMap, groups RoomGroupMap, err error) {
	// NOTE: roomテーブルの行数は少ないことを想定しているため、テーブルスキャンをしている。
	{
		names = make(RoomNameMap)
		var rows *sql.Rows
		rows, err = rst.tx.Query(`
			SELECT room_id, name FROM room
		`)
		if err != nil {
			return
		}
		defer rows.Close()
		for rows.Next() {
			var id RoomID
			var name string
			if err = rows.Scan(&id, &name); err != nil {
				return
			}
			names[id] = name
		}
	}

	{
		groups = make(RoomGroupMap)
		var rows *sql.Rows
		rows, err = rst.tx.Query(`
			SELECT building_name, floor, room_id FROM room
			GROUP BY building_name, floor, room_id
		`)
		defer rows.Close()
		for rows.Next() {
			var bname BuildingName
			var floor FloorID
			var id RoomID
			if err = rows.Scan(&bname, &floor, &id); err != nil {
				return
			}

			if _, ok := groups[bname]; !ok {
				groups[bname] = make(map[FloorID][]RoomID)
			}
			if _, ok := groups[bname][floor]; !ok {
				groups[bname][floor] = make([]RoomID, 0)
			}
			groups[bname][floor] = append(groups[bname][floor], id)
		}
	}
	return
}

func (rsm *RoomStatusManager) getSensorStatusFromCache(id RoomID) ([]SensorStatus, bool) {
	rsm.cacheLock.RLock()
	defer rsm.cacheLock.RUnlock()

	cache, ok := rsm.sensorCache[id]
	if ok {
		array := make([]SensorStatus, 0, len(cache))
		for i := range cache {
			if cache[i].expire.After(time.Now()) {
				array = append(array, cache[i])
			}
		}
		return array, len(array) > 0
	}
	return []SensorStatus{}, false
}

// すべてのセンサーの状態をキャッシュする
func (rsm *RoomStatusManager) cacheUpdater(ctx context.Context) {
	log.Println("starting UpdateStatusWorker")

	tick := time.NewTicker(INTERVAL)
	for {
		log.Println("update status")

		for err := range rsm.updateAllSensorStatuses() {
			log.Println(err)
		}

		select {
		case <-ctx.Done():
			return
		case <-tick.C:
		}
	}
}

func (rsm *RoomStatusManager) updateAllSensorStatuses() []error {
	errCh := make(chan error)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()
		tx, err := rsm.db.Begin()
		if err != nil {
			errCh <- err
			return
		}
		defer tx.Rollback()

		rows, err := tx.Query(
			`SELECT room_id, thing_name FROM thing`,
		)
		if err != nil {
			errCh <- err
			return
		}

		for rows.Next() {
			var id RoomID
			var name ThingName
			rows.Scan(&id, &name)

			// start async update
			wg.Add(1)
			go func(id RoomID, name ThingName) {
				defer wg.Done()
				if err := rsm.updateSensorStatus(id, name); err != nil {
					errCh <- err
					return
				}
			}(id, name)
		}
	}()

	go func() {
		wg.Wait()
		close(errCh)
	}()

	errs := []error{}
	for err := range errCh {
		errs = append(errs, err)
	}
	return errs
}

// センサーで測定した部屋の状態を、DBに反映する。
func (rsm *RoomStatusManager) updateSensorStatus(id RoomID, thingName ThingName) error {
	var stat SensorStatus

	prop, err := rsm.thingworx.Properties(thingName)
	if err != nil {
		return err
	}
	stat.Temperature, err = prop.M("temperature").Float64()
	if err != nil {
		return err
	}
	stat.Humidity, err = prop.M("humidity").Float64()
	if err != nil {
		return err
	}
	stat.IsConnected, err = prop.M("isConnected").Bool()
	if err != nil {
		return err
	}

	rsm.cacheLock.Lock()
	defer rsm.cacheLock.Unlock()
	if _, ok := rsm.sensorCache[id]; !ok {
		rsm.sensorCache[id] = map[ThingName]SensorStatus{}
	}
	rsm.sensorCache[id][thingName] = stat
	return nil
}
<commit_msg>fix Scan error<commit_after>package main

import (
	"context"
	"database/sql"
	"log"
	"net/http"
	"sync"
	"time"
)

type RoomNameMap map[RoomID]string
type RoomGroupMap map[BuildingName]map[FloorID][]RoomID

const (
	INTERVAL = 1 * time.Minute
)

type RoomStatus struct {
	RoomID  RoomID         `json:"id"`
	Sensors []SensorStatus `json:"sensors"`

	Hot     uint64 `json:"hot"`
	Comfort uint64 `json:"comfort"`
	Cold    uint64 `json:"cold"`
	lock    sync.RWMutex
}

type MyVote struct {
	Vote      VoteChoice `json:"vote"`
	Timestamp int64      `json:"timestamp"`
}

type RoomStatusManager struct {
	db        *sql.DB
	thingworx *ThingWorxClient

	sensorCache map[RoomID]map[ThingName]SensorStatus
	cacheLock   sync.RWMutex
}

type RoomStatusTx struct {
	rsm *RoomStatusManager
	tx  *sql.Tx
	s   *Session
}

type SensorStatus struct {
	Temperature float64 `json:"temperature"`
	Humidity    float64 `json:"humidity"`
	IsConnected bool    `json:"isConnected"`

	expire time.Time
}

func NewRoomStatusManager(db *sql.DB, thingworx *ThingWorxClient, ctx context.Context) *RoomStatusManager {
	// create RSM
	rs := &RoomStatusManager{}
	rs.db = db
	rs.thingworx = thingworx
	rs.sensorCache = make(map[RoomID]map[ThingName]SensorStatus)

	go rs.cacheUpdater(ctx)
	return rs
}

func (rsm *RoomStatusManager) GetTx(w http.ResponseWriter, req *http.Request) (*RoomStatusTx, error) {
	tx, err := rsm.db.Begin()
	if err != nil {
		return nil, err
	}
	s := GetSession(w, req, tx)
	return &RoomStatusTx{
		rsm: rsm,
		tx:  tx,
		s:   s,
	}, nil
}

func (rst *RoomStatusTx) Rollback() error {
	return rst.tx.Rollback()
}

func (rst *RoomStatusTx) Commit() error {
	return rst.tx.Commit()
}

func (rst *RoomStatusTx) GetRoomName(id RoomID) (name string, err error) {
	err = rst.tx.QueryRow(
		`SELECT name FROM room
		WHERE room_id=?`,
		id,
	).Scan(&name)
	return
}

// 投票内容を取得する。未投票の場合はnilを返す
func (rst *RoomStatusTx) GetMyVote(id RoomID) (vote *MyVote, err error) {
	var v Vote
	if err = rst.tx.QueryRow(
		`SELECT choice, timestamp FROM vote
			WHERE session_id=? AND room_id=?`,
		rst.s.SessionID, id,
	).Scan(&v.Choice, &v.Timestamp); err != nil {
		if err == sql.ErrNoRows {
			// 未投票の状態。
			err = nil
			return
		}
		return
	}

	vote = &MyVote{
		Vote:      v.Choice,
		Timestamp: v.Timestamp.Unix(),
	}
	return vote, err
}

func (rst *RoomStatusTx) GetStatus(id RoomID) (*RoomStatus, error) {
	rs := &RoomStatus{
		RoomID: id,
	}

	var ok bool
	rs.Sensors, ok = rst.rsm.getSensorStatusFromCache(id)
	if !ok {
		// センサーの状態を更新できていない状態。
		rs.Sensors = []SensorStatus{}
	}

	rows, err := rst.tx.Query(
		`SELECT choice, count(vote_id) FROM vote
		WHERE room_id=?
		GROUP BY choice`,
		id,
	)
	if err != nil {
		return nil, err
	}
	for rows.Next() {
		var choice VoteChoice
		var count uint64
		if err := rows.Scan(&choice, &count); err != nil {
			return nil, err
		}
		switch choice {
		case Hot:
			rs.Hot = count
		case Comfort:
			rs.Comfort = count
		case Cold:
			rs.Cold = count
		}
	}
	return rs, nil
}

func (rst *RoomStatusTx) Vote(id RoomID, choice VoteChoice) error {
	vote := Vote{
		RoomID: id,
		S:      rst.s,
	}

	if err := rst.tx.QueryRow(
		`SELECT vote_id FROM vote
		WHERE session_id=? AND room_id=?`,
		rst.s.SessionID, id,
	).Scan(&vote.VoteID); err != nil {
		switch err {
		case sql.ErrNoRows:
			// 未投票であることを表す、0を代入
			vote.VoteID = VoteID(0)
		default:
			return err
		}
	}

	return vote.UpdateChoice(rst.tx, choice)
}

func (rst *RoomStatusTx) GetAllRoomsInfo() (names RoomNameMap, groups RoomGroupMap, err error) {
	// NOTE: roomテーブルの行数は少ないことを想定しているため、テーブルスキャンをしている。
	{
		names = make(RoomNameMap)
		var rows *sql.Rows
		rows, err = rst.tx.Query(`
			SELECT room_id, name FROM room
		`)
		if err != nil {
			return
		}
		defer rows.Close()
		for rows.Next() {
			var id RoomID
			var name string
			if err = rows.Scan(&id, &name); err != nil {
				return
			}
			names[id] = name
		}
	}

	{
		groups = make(RoomGroupMap)
		var rows *sql.Rows
		rows, err = rst.tx.Query(`
			SELECT building_name, floor, room_id FROM room
			GROUP BY building_name, floor, room_id
		`)
		defer rows.Close()
		for rows.Next() {
			var bname BuildingName
			var floor FloorID
			var id RoomID
			if err = rows.Scan((*string)(&bname), &floor, &id); err != nil {
				return
			}

			if _, ok := groups[bname]; !ok {
				groups[bname] = make(map[FloorID][]RoomID)
			}
			if _, ok := groups[bname][floor]; !ok {
				groups[bname][floor] = make([]RoomID, 0)
			}
			groups[bname][floor] = append(groups[bname][floor], id)
		}
	}
	return
}

func (rsm *RoomStatusManager) getSensorStatusFromCache(id RoomID) ([]SensorStatus, bool) {
	rsm.cacheLock.RLock()
	defer rsm.cacheLock.RUnlock()

	cache, ok := rsm.sensorCache[id]
	if ok {
		array := make([]SensorStatus, 0, len(cache))
		for i := range cache {
			if cache[i].expire.After(time.Now()) {
				array = append(array, cache[i])
			}
		}
		return array, len(array) > 0
	}
	return []SensorStatus{}, false
}

// すべてのセンサーの状態をキャッシュする
func (rsm *RoomStatusManager) cacheUpdater(ctx context.Context) {
	log.Println("starting UpdateStatusWorker")

	tick := time.NewTicker(INTERVAL)
	for {
		log.Println("update status")

		for err := range rsm.updateAllSensorStatuses() {
			log.Println(err)
		}

		select {
		case <-ctx.Done():
			return
		case <-tick.C:
		}
	}
}

func (rsm *RoomStatusManager) updateAllSensorStatuses() []error {
	errCh := make(chan error)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()
		tx, err := rsm.db.Begin()
		if err != nil {
			errCh <- err
			return
		}
		defer tx.Rollback()

		rows, err := tx.Query(
			`SELECT room_id, thing_name FROM thing`,
		)
		if err != nil {
			errCh <- err
			return
		}

		for rows.Next() {
			var id RoomID
			var name ThingName
			rows.Scan(&id, &name)

			// start async update
			wg.Add(1)
			go func(id RoomID, name ThingName) {
				defer wg.Done()
				if err := rsm.updateSensorStatus(id, name); err != nil {
					errCh <- err
					return
				}
			}(id, name)
		}
	}()

	go func() {
		wg.Wait()
		close(errCh)
	}()

	errs := []error{}
	for err := range errCh {
		errs = append(errs, err)
	}
	return errs
}

// センサーで測定した部屋の状態を、DBに反映する。
func (rsm *RoomStatusManager) updateSensorStatus(id RoomID, thingName ThingName) error {
	var stat SensorStatus

	prop, err := rsm.thingworx.Properties(thingName)
	if err != nil {
		return err
	}
	stat.Temperature, err = prop.M("temperature").Float64()
	if err != nil {
		return err
	}
	stat.Humidity, err = prop.M("humidity").Float64()
	if err != nil {
		return err
	}
	stat.IsConnected, err = prop.M("isConnected").Bool()
	if err != nil {
		return err
	}

	rsm.cacheLock.Lock()
	defer rsm.cacheLock.Unlock()
	if _, ok := rsm.sensorCache[id]; !ok {
		rsm.sensorCache[id] = map[ThingName]SensorStatus{}
	}
	rsm.sensorCache[id][thingName] = stat
	return nil
}
<|endoftext|>


<commit_before>package crypto

import (
	"io"

	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type Digest interface {
	Verify(io.Reader) error
	VerifyFilePath(filePath string, fs boshsys.FileSystem) error
	Algorithm() Algorithm
	String() string
}

var _ Digest = digestImpl{}

type Algorithm interface {
	CreateDigest(io.Reader) (Digest, error)
	Name() string
}

var _ Algorithm = algorithmSHAImpl{}
var _ Algorithm = unknownAlgorithmImpl{}
<commit_msg>Update crypto interface to ensure MultipleDigest conforms<commit_after>package crypto

import (
	"io"

	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type Digest interface {
	Verify(io.Reader) error
	VerifyFilePath(filePath string, fs boshsys.FileSystem) error
	Algorithm() Algorithm
	String() string
}

var _ Digest = digestImpl{}
var _ Digest = MultipleDigest{}

type Algorithm interface {
	CreateDigest(io.Reader) (Digest, error)
	Name() string
}

var _ Algorithm = algorithmSHAImpl{}
var _ Algorithm = unknownAlgorithmImpl{}
<|endoftext|>


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

package builder

import (
	"crypto/rand"
	"encoding/binary"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil/gcs"
)

const DefaultP = 20

// GCSBuilder is a utility class that makes building GCS filters convenient.
type GCSBuilder struct {
	p    uint8
	key  [gcs.KeySize]byte
	data [][]byte
	err  error
}

// RandomKey is a utility function that returns a cryptographically random
// [gcs.KeySize]byte usable as a key for a GCS filter.
func RandomKey() ([gcs.KeySize]byte, error) {
	var key [gcs.KeySize]byte

	// Read a byte slice from rand.Reader.
	randKey := make([]byte, gcs.KeySize)
	_, err := rand.Read(randKey)

	// This shouldn't happen unless the user is on a system that doesn't
	// have a system CSPRNG. OK to panic in this case.
	if err != nil {
		return key, err
	}

	// Copy the byte slice to a [gcs.KeySize]byte array and return it.
	copy(key[:], randKey[:])
	return key, nil
}

// DeriveKey is a utility function that derives a key from a chainhash.Hash by
// truncating the bytes of the hash to the appopriate key size.
func DeriveKey(keyHash *chainhash.Hash) [gcs.KeySize]byte {
	var key [gcs.KeySize]byte
	copy(key[:], keyHash.CloneBytes()[:])
	return key
}

// OutPointToFilterEntry is a utility function that derives a filter entry from
// a wire.OutPoint in a standardized way for use with both building and querying
// filters.
func OutPointToFilterEntry(outpoint wire.OutPoint) []byte {
	// Size of the hash plus size of int32 index
	data := make([]byte, chainhash.HashSize+4)
	copy(data[:], outpoint.Hash.CloneBytes()[:])
	binary.BigEndian.PutUint32(data[chainhash.HashSize:], outpoint.Index)
	return data
}

// Key retrieves the key with which the builder will build a filter. This is
// useful if the builder is created with a random initial key.
func (b *GCSBuilder) Key() ([gcs.KeySize]byte, error) {
	// Do nothing if the builder's errored out.
	if b.err != nil {
		return [gcs.KeySize]byte{}, b.err
	}

	return b.key, nil
}

// SetKey sets the key with which the builder will build a filter to the passed
// [gcs.KeySize]byte.
func (b *GCSBuilder) SetKey(key [gcs.KeySize]byte) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	copy(b.key[:], key[:])
	return b
}

// SetKeyFromHash sets the key with which the builder will build a filter to a
// key derived from the passed chainhash.Hash using DeriveKey().
func (b *GCSBuilder) SetKeyFromHash(keyHash *chainhash.Hash) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	return b.SetKey(DeriveKey(keyHash))
}

// SetP sets the filter's probability after calling Builder().
func (b *GCSBuilder) SetP(p uint8) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	// Basic sanity check.
	if p > 32 {
		b.err = gcs.ErrPTooBig
		return b
	}

	b.p = p
	return b
}

// Preallocate sets the estimated filter size after calling Builder() to reduce
// the probability of memory reallocations. If the builder has already had data
// added to it, SetN has no effect.
func (b *GCSBuilder) Preallocate(n uint32) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	if len(b.data) == 0 {
		b.data = make([][]byte, 0, n)
	}
	return b
}

// AddEntry adds a []byte to the list of entries to be included in the GCS
// filter when it's built.
func (b *GCSBuilder) AddEntry(data []byte) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	b.data = append(b.data, data)
	return b
}

// AddEntries adds all the []byte entries in a [][]byte to the list of entries
// to be included in the GCS filter when it's built.
func (b *GCSBuilder) AddEntries(data [][]byte) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	for _, entry := range data {
		b.AddEntry(entry)
	}
	return b
}

// AddOutPoint adds a wire.OutPoint to the list of entries to be included in the
// GCS filter when it's built.
func (b *GCSBuilder) AddOutPoint(outpoint wire.OutPoint) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	return b.AddEntry(OutPointToFilterEntry(outpoint))
}

// AddHash adds a chainhash.Hash to the list of entries to be included in the
// GCS filter when it's built.
func (b *GCSBuilder) AddHash(hash *chainhash.Hash) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	return b.AddEntry(hash.CloneBytes())
}

// AddScript adds all the data pushed in the script serialized as the passed
// []byte to the list of entries to be included in the GCS filter when it's
// built. T
func (b *GCSBuilder) AddScript(script []byte) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	data, err := txscript.PushedData(script)
	if err != nil {
		b.err = err
		return b
	}
	return b.AddEntries(data)
}

// Build returns a function which builds a GCS filter with the given parameters
// and data.
func (b *GCSBuilder) Build() (*gcs.Filter, error) {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return nil, b.err
	}

	return gcs.BuildGCSFilter(b.p, b.key, b.data)
}

// WithKeyPN creates a GCSBuilder with specified key and the passed
// probability and estimated filter size.
func WithKeyPN(key [gcs.KeySize]byte, p uint8, n uint32) *GCSBuilder {
	b := GCSBuilder{}
	return b.SetKey(key).SetP(p).Preallocate(n)
}

// WithKeyP creates a GCSBuilder with specified key and the passed
// probability. Estimated filter size is set to zero, which means more
// reallocations are done when building the filter.
func WithKeyP(key [gcs.KeySize]byte, p uint8) *GCSBuilder {
	return WithKeyPN(key, p, 0)
}

// WithKey creates a GCSBuilder with specified key. Probability is set to
// 20 (2^-20 collision probability). Estimated filter size is set to zero, which
// means more reallocations are done when building the filter.
func WithKey(key [gcs.KeySize]byte) *GCSBuilder {
	return WithKeyPN(key, DefaultP, 0)
}

// WithKeyHashPN creates a GCSBuilder with key derived from the specified
// chainhash.Hash and the passed probability and estimated filter size.
func WithKeyHashPN(keyHash *chainhash.Hash, p uint8, n uint32) *GCSBuilder {
	return WithKeyPN(DeriveKey(keyHash), p, n)
}

// WithKeyHashP creates a GCSBuilder with key derived from the specified
// chainhash.Hash and the passed probability. Estimated filter size is set to
// zero, which means more reallocations are done when building the filter.
func WithKeyHashP(keyHash *chainhash.Hash, p uint8) *GCSBuilder {
	return WithKeyHashPN(keyHash, p, 0)
}

// WithKeyHash creates a GCSBuilder with key derived from the specified
// chainhash.Hash. Probability is set to 20 (2^-20 collision probability).
// Estimated filter size is set to zero, which means more reallocations are
// done when building the filter.
func WithKeyHash(keyHash *chainhash.Hash) *GCSBuilder {
	return WithKeyHashPN(keyHash, DefaultP, 0)
}

// WithRandomKeyPN creates a GCSBuilder with a cryptographically random
// key and the passed probability and estimated filter size.
func WithRandomKeyPN(p uint8, n uint32) *GCSBuilder {
	key, err := RandomKey()
	if err != nil {
		b := GCSBuilder{err: err}
		return &b
	}
	return WithKeyPN(key, p, n)
}

// WithRandomKeyP creates a GCSBuilder with a cryptographically random
// key and the passed probability. Estimated filter size is set to zero, which
// means more reallocations are done when building the filter.
func WithRandomKeyP(p uint8) *GCSBuilder {
	return WithRandomKeyPN(p, 0)
}

// WithRandomKey creates a GCSBuilder with a cryptographically random
// key. Probability is set to 20 (2^-20 collision probability). Estimated
// filter size is set to zero, which means more reallocations are done when
// building the filter.
func WithRandomKey() *GCSBuilder {
	return WithRandomKeyPN(DefaultP, 0)
}
<commit_msg>Add comment to DefaultP to fix lint issue with exported constant.<commit_after>// Copyright (c) 2017 The btcsuite developers
// Copyright (c) 2017 The Lightning Network Developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package builder

import (
	"crypto/rand"
	"encoding/binary"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil/gcs"
)

// DefaultP is the default collision probability (2^-20)
const DefaultP = 20

// GCSBuilder is a utility class that makes building GCS filters convenient.
type GCSBuilder struct {
	p    uint8
	key  [gcs.KeySize]byte
	data [][]byte
	err  error
}

// RandomKey is a utility function that returns a cryptographically random
// [gcs.KeySize]byte usable as a key for a GCS filter.
func RandomKey() ([gcs.KeySize]byte, error) {
	var key [gcs.KeySize]byte

	// Read a byte slice from rand.Reader.
	randKey := make([]byte, gcs.KeySize)
	_, err := rand.Read(randKey)

	// This shouldn't happen unless the user is on a system that doesn't
	// have a system CSPRNG. OK to panic in this case.
	if err != nil {
		return key, err
	}

	// Copy the byte slice to a [gcs.KeySize]byte array and return it.
	copy(key[:], randKey[:])
	return key, nil
}

// DeriveKey is a utility function that derives a key from a chainhash.Hash by
// truncating the bytes of the hash to the appopriate key size.
func DeriveKey(keyHash *chainhash.Hash) [gcs.KeySize]byte {
	var key [gcs.KeySize]byte
	copy(key[:], keyHash.CloneBytes()[:])
	return key
}

// OutPointToFilterEntry is a utility function that derives a filter entry from
// a wire.OutPoint in a standardized way for use with both building and querying
// filters.
func OutPointToFilterEntry(outpoint wire.OutPoint) []byte {
	// Size of the hash plus size of int32 index
	data := make([]byte, chainhash.HashSize+4)
	copy(data[:], outpoint.Hash.CloneBytes()[:])
	binary.BigEndian.PutUint32(data[chainhash.HashSize:], outpoint.Index)
	return data
}

// Key retrieves the key with which the builder will build a filter. This is
// useful if the builder is created with a random initial key.
func (b *GCSBuilder) Key() ([gcs.KeySize]byte, error) {
	// Do nothing if the builder's errored out.
	if b.err != nil {
		return [gcs.KeySize]byte{}, b.err
	}

	return b.key, nil
}

// SetKey sets the key with which the builder will build a filter to the passed
// [gcs.KeySize]byte.
func (b *GCSBuilder) SetKey(key [gcs.KeySize]byte) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	copy(b.key[:], key[:])
	return b
}

// SetKeyFromHash sets the key with which the builder will build a filter to a
// key derived from the passed chainhash.Hash using DeriveKey().
func (b *GCSBuilder) SetKeyFromHash(keyHash *chainhash.Hash) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	return b.SetKey(DeriveKey(keyHash))
}

// SetP sets the filter's probability after calling Builder().
func (b *GCSBuilder) SetP(p uint8) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	// Basic sanity check.
	if p > 32 {
		b.err = gcs.ErrPTooBig
		return b
	}

	b.p = p
	return b
}

// Preallocate sets the estimated filter size after calling Builder() to reduce
// the probability of memory reallocations. If the builder has already had data
// added to it, SetN has no effect.
func (b *GCSBuilder) Preallocate(n uint32) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	if len(b.data) == 0 {
		b.data = make([][]byte, 0, n)
	}
	return b
}

// AddEntry adds a []byte to the list of entries to be included in the GCS
// filter when it's built.
func (b *GCSBuilder) AddEntry(data []byte) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	b.data = append(b.data, data)
	return b
}

// AddEntries adds all the []byte entries in a [][]byte to the list of entries
// to be included in the GCS filter when it's built.
func (b *GCSBuilder) AddEntries(data [][]byte) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	for _, entry := range data {
		b.AddEntry(entry)
	}
	return b
}

// AddOutPoint adds a wire.OutPoint to the list of entries to be included in the
// GCS filter when it's built.
func (b *GCSBuilder) AddOutPoint(outpoint wire.OutPoint) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	return b.AddEntry(OutPointToFilterEntry(outpoint))
}

// AddHash adds a chainhash.Hash to the list of entries to be included in the
// GCS filter when it's built.
func (b *GCSBuilder) AddHash(hash *chainhash.Hash) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	return b.AddEntry(hash.CloneBytes())
}

// AddScript adds all the data pushed in the script serialized as the passed
// []byte to the list of entries to be included in the GCS filter when it's
// built. T
func (b *GCSBuilder) AddScript(script []byte) *GCSBuilder {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return b
	}

	data, err := txscript.PushedData(script)
	if err != nil {
		b.err = err
		return b
	}
	return b.AddEntries(data)
}

// Build returns a function which builds a GCS filter with the given parameters
// and data.
func (b *GCSBuilder) Build() (*gcs.Filter, error) {
	// Do nothing if the builder's already errored out.
	if b.err != nil {
		return nil, b.err
	}

	return gcs.BuildGCSFilter(b.p, b.key, b.data)
}

// WithKeyPN creates a GCSBuilder with specified key and the passed
// probability and estimated filter size.
func WithKeyPN(key [gcs.KeySize]byte, p uint8, n uint32) *GCSBuilder {
	b := GCSBuilder{}
	return b.SetKey(key).SetP(p).Preallocate(n)
}

// WithKeyP creates a GCSBuilder with specified key and the passed
// probability. Estimated filter size is set to zero, which means more
// reallocations are done when building the filter.
func WithKeyP(key [gcs.KeySize]byte, p uint8) *GCSBuilder {
	return WithKeyPN(key, p, 0)
}

// WithKey creates a GCSBuilder with specified key. Probability is set to
// 20 (2^-20 collision probability). Estimated filter size is set to zero, which
// means more reallocations are done when building the filter.
func WithKey(key [gcs.KeySize]byte) *GCSBuilder {
	return WithKeyPN(key, DefaultP, 0)
}

// WithKeyHashPN creates a GCSBuilder with key derived from the specified
// chainhash.Hash and the passed probability and estimated filter size.
func WithKeyHashPN(keyHash *chainhash.Hash, p uint8, n uint32) *GCSBuilder {
	return WithKeyPN(DeriveKey(keyHash), p, n)
}

// WithKeyHashP creates a GCSBuilder with key derived from the specified
// chainhash.Hash and the passed probability. Estimated filter size is set to
// zero, which means more reallocations are done when building the filter.
func WithKeyHashP(keyHash *chainhash.Hash, p uint8) *GCSBuilder {
	return WithKeyHashPN(keyHash, p, 0)
}

// WithKeyHash creates a GCSBuilder with key derived from the specified
// chainhash.Hash. Probability is set to 20 (2^-20 collision probability).
// Estimated filter size is set to zero, which means more reallocations are
// done when building the filter.
func WithKeyHash(keyHash *chainhash.Hash) *GCSBuilder {
	return WithKeyHashPN(keyHash, DefaultP, 0)
}

// WithRandomKeyPN creates a GCSBuilder with a cryptographically random
// key and the passed probability and estimated filter size.
func WithRandomKeyPN(p uint8, n uint32) *GCSBuilder {
	key, err := RandomKey()
	if err != nil {
		b := GCSBuilder{err: err}
		return &b
	}
	return WithKeyPN(key, p, n)
}

// WithRandomKeyP creates a GCSBuilder with a cryptographically random
// key and the passed probability. Estimated filter size is set to zero, which
// means more reallocations are done when building the filter.
func WithRandomKeyP(p uint8) *GCSBuilder {
	return WithRandomKeyPN(p, 0)
}

// WithRandomKey creates a GCSBuilder with a cryptographically random
// key. Probability is set to 20 (2^-20 collision probability). Estimated
// filter size is set to zero, which means more reallocations are done when
// building the filter.
func WithRandomKey() *GCSBuilder {
	return WithRandomKeyPN(DefaultP, 0)
}
<|endoftext|>


<commit_before>// Copyright 2016 Marcel Gotsch. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package goserv

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestRouteHandlerChain(t *testing.T) {
	w := httptest.NewRecorder()
	res := &responseWriter{w: w}
	req, _ := http.NewRequest(http.MethodGet, "/", nil)
	history := newHistoryHandler()

	createRequestContext(req)
	ctx := Context(req)

	route := newRoute("/", false, false)

	route.All(history.Handler("1"))
	route.All(history.Handler("2"))
	route.Get(history.WriteHandler("3"))
	route.Get(history.Handler("4"))
	route.Put(history.Handler("5"))

	route.serveHTTP(res, req)
	if err := ctx.err; err != nil {
		t.Errorf("Serve error: %v", err)
	}

	if history.Len() != 3 {
		t.Fatalf("Wrong write count: %d != 3", history.Len())
	}

	for index, value := range []string{"1", "2", "3"} {
		if history.At(index) != value {
			t.Errorf("Invalid write: %s != %s", history.At(index), value)
		}
	}

	if res.Code() != http.StatusOK {
		t.Errorf("Wrong status code: %d != %d", res.Code(), http.StatusOK)
	}

	if w.Body.String() != "3" {
		t.Errorf("Wrong body content: %s != %s", w.Body.String(), "3")
	}
}

func TestRouteRest(t *testing.T) {
	w := httptest.NewRecorder()
	res := &responseWriter{w: w}
	req, _ := http.NewRequest("", "/", nil)
	history := newHistoryHandler()

	createRequestContext(req)

	route := newRoute("/", false, false)

	route.Get(history.WriteHandler("get-handler"))
	route.Rest(history.WriteHandler("rest-handler"))

	for _, method := range methodNames {
		req.Method = method

		route.serveHTTP(res, req)

		wanted := "rest-handler"
		if method == http.MethodGet {
			wanted = "get-handler"
		}

		if first := history.At(0); first != wanted {
			t.Errorf("Wrong write value, wanted: %q, got: %q", "get-handler", wanted, first)
		}

		history.Clear()
	}
}
<commit_msg>Fix arguments in route_test<commit_after>// Copyright 2016 Marcel Gotsch. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package goserv

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestRouteHandlerChain(t *testing.T) {
	w := httptest.NewRecorder()
	res := &responseWriter{w: w}
	req, _ := http.NewRequest(http.MethodGet, "/", nil)
	history := newHistoryHandler()

	createRequestContext(req)
	ctx := Context(req)

	route := newRoute("/", false, false)

	route.All(history.Handler("1"))
	route.All(history.Handler("2"))
	route.Get(history.WriteHandler("3"))
	route.Get(history.Handler("4"))
	route.Put(history.Handler("5"))

	route.serveHTTP(res, req)
	if err := ctx.err; err != nil {
		t.Errorf("Serve error: %v", err)
	}

	if history.Len() != 3 {
		t.Fatalf("Wrong write count: %d != 3", history.Len())
	}

	for index, value := range []string{"1", "2", "3"} {
		if history.At(index) != value {
			t.Errorf("Invalid write: %s != %s", history.At(index), value)
		}
	}

	if res.Code() != http.StatusOK {
		t.Errorf("Wrong status code: %d != %d", res.Code(), http.StatusOK)
	}

	if w.Body.String() != "3" {
		t.Errorf("Wrong body content: %s != %s", w.Body.String(), "3")
	}
}

func TestRouteRest(t *testing.T) {
	w := httptest.NewRecorder()
	res := &responseWriter{w: w}
	req, _ := http.NewRequest("", "/", nil)
	history := newHistoryHandler()

	createRequestContext(req)

	route := newRoute("/", false, false)

	route.Get(history.WriteHandler("get-handler"))
	route.Rest(history.WriteHandler("rest-handler"))

	for _, method := range methodNames {
		req.Method = method

		route.serveHTTP(res, req)

		wanted := "rest-handler"
		if method == http.MethodGet {
			wanted = "get-handler"
		}

		if first := history.At(0); first != wanted {
			t.Errorf("Wrong write value, wanted: %q, got: %q", wanted, first)
		}

		history.Clear()
	}
}
<|endoftext|>


<commit_before>package rpc

import (
	"errors"
	"io"
	"net"

	"github.com/Sirupsen/logrus"
)

type Client struct {
	end       chan struct{}
	requests  chan *Message
	send      chan *Message
	responses chan *Message
	seq       uint32
	messages  map[uint32]*Message
	wire      *Wire
}

func NewClient(conn net.Conn) *Client {
	c := &Client{
		wire:      NewWire(conn),
		end:       make(chan struct{}, 1024),
		requests:  make(chan *Message, 1024),
		send:      make(chan *Message, 1024),
		responses: make(chan *Message, 1024),
		messages:  map[uint32]*Message{},
	}
	go c.loop()
	go c.write()
	go c.read()
	return c
}

func (c *Client) WriteAt(buf []byte, offset int64) (int, error) {
	return c.operation(TypeWrite, buf, offset)
}

func (c *Client) ReadAt(buf []byte, offset int64) (int, error) {
	return c.operation(TypeRead, buf, offset)
}

func (c *Client) operation(op uint32, buf []byte, offset int64) (int, error) {
	msg := Message{
		Type:   op,
		Offset: offset,
		Data:   buf,
	}
	msg.Add(1)
	c.requests <- &msg
	msg.Wait()

	if msg.Type == TypeError {
		return 0, errors.New(string(msg.Data))
	}

	if msg.Type == TypeEOF {
		return len(msg.Data), io.EOF
	}

	return len(msg.Data), nil
}

func (c *Client) Close() {
	c.end <- struct{}{}
	close(c.send)
}

func (c *Client) loop() {
	for {
		select {
		case <-c.end:
			break
		case req := <-c.requests:
			c.handleRequest(req)
		case resp := <-c.responses:
			c.handleResponse(resp)
		}
	}
}

func (c *Client) nextSeq() uint32 {
	c.seq++
	return c.seq
}

func (c *Client) handleRequest(req *Message) {
	req.Seq = c.nextSeq()
	c.messages[req.Seq] = req
	c.send <- req
}

func (c *Client) handleResponse(resp *Message) {
	if req, ok := c.messages[resp.Seq]; ok {
		// can probably optimize away this copy
		if len(resp.Data) > 0 {
			copy(req.Data, resp.Data)
		}
		req.Type = resp.Type
		req.Done()
	}
}

func (c *Client) write() {
	for msg := range c.send {
		if err := c.wire.Write(msg); err != nil {
			c.responses <- &Message{
				Seq:  msg.Seq,
				Type: TypeError,
				Data: []byte(err.Error()),
			}
		}
	}
}

func (c *Client) read() {
	for {
		msg, err := c.wire.Read()
		if err != nil {
			logrus.Errorf("Error reading from wire: %v", err)
		}
		c.responses <- msg
	}
}
<commit_msg>Fix memory issue<commit_after>package rpc

import (
	"errors"
	"io"
	"net"

	"github.com/Sirupsen/logrus"
)

type Client struct {
	end       chan struct{}
	requests  chan *Message
	send      chan *Message
	responses chan *Message
	seq       uint32
	messages  map[uint32]*Message
	wire      *Wire
}

func NewClient(conn net.Conn) *Client {
	c := &Client{
		wire:      NewWire(conn),
		end:       make(chan struct{}, 1024),
		requests:  make(chan *Message, 1024),
		send:      make(chan *Message, 1024),
		responses: make(chan *Message, 1024),
		messages:  map[uint32]*Message{},
	}
	go c.loop()
	go c.write()
	go c.read()
	return c
}

func (c *Client) WriteAt(buf []byte, offset int64) (int, error) {
	return c.operation(TypeWrite, buf, offset)
}

func (c *Client) ReadAt(buf []byte, offset int64) (int, error) {
	return c.operation(TypeRead, buf, offset)
}

func (c *Client) operation(op uint32, buf []byte, offset int64) (int, error) {
	msg := Message{
		Type:   op,
		Offset: offset,
		Data:   buf,
	}
	msg.Add(1)
	c.requests <- &msg
	msg.Wait()

	if msg.Type == TypeError {
		return 0, errors.New(string(msg.Data))
	}

	if msg.Type == TypeEOF {
		return len(msg.Data), io.EOF
	}

	return len(msg.Data), nil
}

func (c *Client) Close() {
	c.end <- struct{}{}
	close(c.send)
}

func (c *Client) loop() {
	for {
		select {
		case <-c.end:
			break
		case req := <-c.requests:
			c.handleRequest(req)
		case resp := <-c.responses:
			c.handleResponse(resp)
		}
	}
}

func (c *Client) nextSeq() uint32 {
	c.seq++
	return c.seq
}

func (c *Client) handleRequest(req *Message) {
	req.Seq = c.nextSeq()
	c.messages[req.Seq] = req
	c.send <- req
}

func (c *Client) handleResponse(resp *Message) {
	if req, ok := c.messages[resp.Seq]; ok {
		delete(c.messages, resp.Seq)
		// can probably optimize away this copy
		if len(resp.Data) > 0 {
			copy(req.Data, resp.Data)
		}
		req.Type = resp.Type
		req.Done()
	}
}

func (c *Client) write() {
	for msg := range c.send {
		if err := c.wire.Write(msg); err != nil {
			c.responses <- &Message{
				Seq:  msg.Seq,
				Type: TypeError,
				Data: []byte(err.Error()),
			}
		}
	}
}

func (c *Client) read() {
	for {
		msg, err := c.wire.Read()
		if err != nil {
			logrus.Errorf("Error reading from wire: %v", err)
		}
		c.responses <- msg
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

package rpc

import "encoding/json"

// NewInProcRPCClient creates an in-process buffer stream attachment to a given
// RPC server.
func NewInProcRPCClient(handler *Server) Client {
	buffer := &inprocBuffer{
		requests:  make(chan []byte, 16),
		responses: make(chan []byte, 16),
	}
	client := &inProcClient{
		server: handler,
		buffer: buffer,
	}
	go handler.ServeCodec(NewJSONCodec(client.buffer))
	return client
}

// inProcClient is an in-process buffer stream attached to an RPC server.
type inProcClient struct {
	server *Server
	buffer *inprocBuffer
}

// Close tears down the request channel of the in-proc client.
func (c *inProcClient) Close() {
	c.buffer.Close()
}

// Send marshals a message into a json format and injects in into the client
// request channel.
func (c *inProcClient) Send(msg interface{}) error {
	d, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	c.buffer.requests <- d
	return nil
}

// Recv reads a message from the response channel and tries to parse it into the
// given msg interface.
func (c *inProcClient) Recv(msg interface{}) error {
	data := <-c.buffer.responses
	return json.Unmarshal(data, &msg)
}

// Returns the collection of modules the RPC server offers.
func (c *inProcClient) SupportedModules() (map[string]string, error) {
	return SupportedModules(c)
}

// inprocBuffer represents the connection between the RPC server and console
type inprocBuffer struct {
	readBuf   []byte      // store remaining request bytes after a partial read
	requests  chan []byte // list with raw serialized requests
	responses chan []byte // list with raw serialized responses
}

// Read will read the next request in json format.
func (b *inprocBuffer) Read(p []byte) (int, error) {
	// last read didn't read entire request, return remaining bytes
	if len(b.readBuf) > 0 {
		n := copy(p, b.readBuf)
		if n < len(b.readBuf) {
			b.readBuf = b.readBuf[:n]
		} else {
			b.readBuf = b.readBuf[:0]
		}
		return n, nil
	}
	// read next request
	req := <-b.requests
	n := copy(p, req)
	if n < len(req) {
		// inprocBuffer too small, store remaining chunk for next read
		b.readBuf = req[n:]
	}
	return n, nil
}

// Write sends the given buffer to the backend.
func (b *inprocBuffer) Write(p []byte) (n int, err error) {
	b.responses <- p
	return len(p), nil
}

// Close cleans up obtained resources.
func (b *inprocBuffer) Close() error {
	close(b.requests)
	close(b.responses)

	return nil
}
<commit_msg>rpc: simplify inproc client<commit_after>// Copyright 2016 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

package rpc

import (
	"encoding/json"
	"io"
	"net"
)

// inProcClient is an in-process buffer stream attached to an RPC server.
type inProcClient struct {
	server *Server
	cl     io.Closer
	enc    *json.Encoder
	dec    *json.Decoder
}

// Close tears down the request channel of the in-proc client.
func (c *inProcClient) Close() {
	c.cl.Close()
}

// NewInProcRPCClient creates an in-process buffer stream attachment to a given
// RPC server.
func NewInProcRPCClient(handler *Server) Client {
	p1, p2 := net.Pipe()
	go handler.ServeCodec(NewJSONCodec(p1))
	return &inProcClient{handler, p2, json.NewEncoder(p2), json.NewDecoder(p2)}
}

// Send marshals a message into a json format and injects in into the client
// request channel.
func (c *inProcClient) Send(msg interface{}) error {
	return c.enc.Encode(msg)
}

// Recv reads a message from the response channel and tries to parse it into the
// given msg interface.
func (c *inProcClient) Recv(msg interface{}) error {
	return c.dec.Decode(msg)
}

// Returns the collection of modules the RPC server offers.
func (c *inProcClient) SupportedModules() (map[string]string, error) {
	return SupportedModules(c)
}
<|endoftext|>


<commit_before>package meritop

// framework_test tests basic workflows of framework impl.
// It uses a scenario with two nodes: 0 as parent, 1 as child.
// The basic idea is that when parent tries to talk to child and vice versa,
// there will be some data transferring and captured by application task.
// Here we have implemented a helper user task to capture those data, test if
// it's passed from framework correctly and unmodified.

import (
	"bytes"
	"fmt"
	"net"
	"testing"
)

func TestFrameworkFlagMetaReady(t *testing.T) {
	m := mustNewMember(t, "framework_test")
	m.Launch()
	defer m.Terminate(t)
	url := fmt.Sprintf("http://%s", m.ClientListeners[0].Addr().String())

	pDataChan := make(chan *tDataBundle, 1)
	cDataChan := make(chan *tDataBundle, 1)
	// simulate two tasks on two nodes -- 0 and 1
	// 0 is parent, 1 is child
	f0 := &framework{
		name:     "framework_test_flagmetaready",
		etcdURLs: []string{url},
		taskID:   0,
		task: &testableTask{
			dataChan: cDataChan,
		},
		topology: NewTreeTopology(2, 1),
		ln:       createListener(t),
	}
	f1 := &framework{
		name:     "framework_test_flagmetaready",
		etcdURLs: []string{url},
		taskID:   1,
		task: &testableTask{
			dataChan: pDataChan,
		},
		topology: NewTreeTopology(2, 1),
		ln:       createListener(t),
	}
	f0.start()
	defer f0.stop()
	f1.start()
	defer f1.stop()

	tests := []struct {
		cMeta string
		pMeta string
	}{
		{"parent", "child"},
		{"ParamReady", "GradientReady"},
	}

	for i, tt := range tests {
		// 0: F#FlagChildMetaReady -> 1: T#ParentMetaReady
		f0.FlagChildMetaReady(tt.cMeta)
		// from child(1)'s view
		data := <-pDataChan
		if data.id != 0 {
			t.Errorf("#%d: parentID want = 0, get = %d", data.id)
		}
		if data.meta != tt.cMeta {
			t.Errorf("#%d: meta want = %s, get = %s", i, tt.cMeta, data.meta)
		}

		// 1: F#FlagParentMetaReady -> 0: T#ChildMetaReady
		f1.FlagParentMetaReady(tt.pMeta)
		// from parent(0)'s view
		data = <-cDataChan
		if data.id != 1 {
			t.Errorf("#%d: parentID want = 1, get = %d", data.id)
		}
		if data.meta != tt.pMeta {
			t.Errorf("#%d: meta want = %s, get = %s", i, tt.pMeta, data.meta)
		}
	}
}

func TestFrameworkDataRequest(t *testing.T) {
	tests := []struct {
		req  string
		resp []byte
	}{
		{"request", []byte("response")},
		{"parameters", []byte{1, 2, 3}},
		{"gradient", []byte{4, 5, 6}},
	}

	dataMap := make(map[string][]byte)
	for _, tt := range tests {
		dataMap[tt.req] = tt.resp
	}

	m := mustNewMember(t, "framework_test")
	m.Launch()
	defer m.Terminate(t)
	url := fmt.Sprintf("http://%s", m.ClientListeners[0].Addr().String())
	l0 := createListener(t)
	l1 := createListener(t)
	addressMap := map[uint64]string{
		0: l0.Addr().String(),
		1: l1.Addr().String(),
	}

	pDataChan := make(chan *tDataBundle, 1)
	cDataChan := make(chan *tDataBundle, 1)
	// simulate two tasks on two nodes -- 0 and 1
	// 0 is parent, 1 is child
	f0 := &framework{
		name:     "framework_test_datarequest",
		etcdURLs: []string{url},
		taskID:   0,
		task: &testableTask{
			dataMap:  dataMap,
			dataChan: cDataChan,
		},
		topology:   NewTreeTopology(2, 1),
		ln:         l0,
		addressMap: addressMap,
	}
	f1 := &framework{
		name:     "framework_test_datarequest",
		etcdURLs: []string{url},
		taskID:   1,
		task: &testableTask{
			dataMap:  dataMap,
			dataChan: pDataChan,
		},
		topology:   NewTreeTopology(2, 1),
		ln:         l1,
		addressMap: addressMap,
	}
	f0.start()
	defer f0.stop()
	f1.start()
	defer f1.stop()

	for i, tt := range tests {
		// 0: F#DataRequest -> 1: T#ServeAsChild -> 0: T#ChildDataReady
		f0.DataRequest(1, tt.req)
		// from child(1)'s view at 1: T#ServeAsChild
		data := <-pDataChan
		if data.id != 0 {
			t.Errorf("#%d: fromID want = 0, get = %d", i, data.id)
		}
		if data.req != tt.req {
			t.Errorf("#%d: req want = %s, get = %s", i, tt.req, data.req)
		}
		// from parent(0)'s view at 0: T#ChildDataReady
		data = <-cDataChan
		if data.id != 1 {
			t.Errorf("#%d: fromID want = 1, get = %d", i, data.id)
		}
		if data.req != tt.req {
			t.Errorf("#%d: req want = %s, get = %s", i, tt.req, data.req)
		}
		if bytes.Compare(data.resp, tt.resp) != 0 {
			t.Errorf("#%d: resp want = %v, get = %v", i, tt.resp, data.resp)
		}

		// 1: F#DataRequest -> 0: T#ServeAsParent -> 1: T#ParentDataReady
		f1.DataRequest(0, tt.req)
		// from parent(0)'s view at 0: T#ServeAsParent
		data = <-cDataChan
		if data.id != 1 {
			t.Errorf("#%d: fromID want = 1, get = %d", i, data.id)
		}
		if data.req != tt.req {
			t.Errorf("#%d: req want = %s, get = %s", i, tt.req, data.req)
		}
		// from child(1)'s view at 1: T#ParentDataReady
		data = <-pDataChan
		if data.id != 0 {
			t.Errorf("#%d: fromID want = 1, get = %d", i, data.id)
		}
		if data.req != tt.req {
			t.Errorf("#%d: req want = %s, get = %s", i, tt.req, data.req)
		}
		if bytes.Compare(data.resp, tt.resp) != 0 {
			t.Errorf("#%d: resp want = %v, get = %v", i, tt.resp, data.resp)
		}
	}
}

type tDataBundle struct {
	id   uint64
	meta string
	req  string
	resp []byte
}

type testableTask struct {
	id        uint64
	framework Framework
	// dataMap will be used to serve data according to request
	dataMap map[string][]byte

	// This channel is used to convey data passed from framework back to the main
	// thread, for checking. Thus it's initialized and passed in from outside.
	//
	// The basic idea is that there are only two nodes -- one parent and one child.
	// When this channel is for parent, it passes information from child.
	dataChan chan *tDataBundle
}

func (t *testableTask) Init(taskID uint64, framework Framework, config Config) {
	t.id = taskID
	t.framework = framework
}
func (t *testableTask) Exit()                 {}
func (t *testableTask) SetEpoch(epoch uint64) {}

func (t *testableTask) ParentMetaReady(fromID uint64, meta string) {
	if t.dataChan != nil {
		t.dataChan <- &tDataBundle{fromID, meta, "", nil}
	}
}

func (t *testableTask) ChildMetaReady(fromID uint64, meta string) {
	t.ParentMetaReady(fromID, meta)
}

func (t *testableTask) ServeAsParent(fromID uint64, req string) []byte {
	if t.dataChan != nil {
		t.dataChan <- &tDataBundle{fromID, "", req, nil}
	}
	return t.dataMap[req]
}
func (t *testableTask) ServeAsChild(fromID uint64, req string) []byte {
	return t.ServeAsParent(fromID, req)
}
func (t *testableTask) ParentDataReady(fromID uint64, req string, resp []byte) {
	if t.dataChan != nil {
		t.dataChan <- &tDataBundle{fromID, "", req, resp}
	}
}

func (t *testableTask) ChildDataReady(fromID uint64, req string, resp []byte) {
	t.ParentDataReady(fromID, req, resp)
}

func createListener(t *testing.T) net.Listener {
	l, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.Listen(\"tcp4\", \"\") failed: %v", err)
	}
	return l
}
<commit_msg>refactor all test checking to DataBundle<commit_after>package meritop

// framework_test tests basic workflows of framework impl.
// It uses a scenario with two nodes: 0 as parent, 1 as child.
// The basic idea is that when parent tries to talk to child and vice versa,
// there will be some data transferring and captured by application task.
// Here we have implemented a helper user task to capture those data, test if
// it's passed from framework correctly and unmodified.

import (
	"fmt"
	"net"
	"reflect"
	"testing"
)

func TestFrameworkFlagMetaReady(t *testing.T) {
	m := mustNewMember(t, "framework_test")
	m.Launch()
	defer m.Terminate(t)
	url := fmt.Sprintf("http://%s", m.ClientListeners[0].Addr().String())

	pDataChan := make(chan *tDataBundle, 1)
	cDataChan := make(chan *tDataBundle, 1)
	// simulate two tasks on two nodes -- 0 and 1
	// 0 is parent, 1 is child
	f0 := &framework{
		name:     "framework_test_flagmetaready",
		etcdURLs: []string{url},
		taskID:   0,
		task: &testableTask{
			dataChan: cDataChan,
		},
		topology: NewTreeTopology(2, 1),
		ln:       createListener(t),
	}
	f1 := &framework{
		name:     "framework_test_flagmetaready",
		etcdURLs: []string{url},
		taskID:   1,
		task: &testableTask{
			dataChan: pDataChan,
		},
		topology: NewTreeTopology(2, 1),
		ln:       createListener(t),
	}
	f0.start()
	defer f0.stop()
	f1.start()
	defer f1.stop()

	tests := []struct {
		cMeta string
		pMeta string
	}{
		{"parent", "child"},
		{"ParamReady", "GradientReady"},
	}

	for i, tt := range tests {
		// 0: F#FlagChildMetaReady -> 1: T#ParentMetaReady
		f0.FlagChildMetaReady(tt.cMeta)
		// from child(1)'s view
		data := <-pDataChan
		expected := &tDataBundle{0, tt.cMeta, "", nil}
		if !reflect.DeepEqual(data, expected) {
			t.Errorf("#%d: data bundle want = %v, get = %v", i, expected, data)
		}

		// 1: F#FlagParentMetaReady -> 0: T#ChildMetaReady
		f1.FlagParentMetaReady(tt.pMeta)
		// from parent(0)'s view
		data = <-cDataChan
		expected = &tDataBundle{1, tt.pMeta, "", nil}
		if !reflect.DeepEqual(data, expected) {
			t.Errorf("#%d: data bundle want = %v, get = %v", i, expected, data)
		}
	}
}

func TestFrameworkDataRequest(t *testing.T) {
	tests := []struct {
		req  string
		resp []byte
	}{
		{"request", []byte("response")},
		{"parameters", []byte{1, 2, 3}},
		{"gradient", []byte{4, 5, 6}},
	}

	dataMap := make(map[string][]byte)
	for _, tt := range tests {
		dataMap[tt.req] = tt.resp
	}

	m := mustNewMember(t, "framework_test")
	m.Launch()
	defer m.Terminate(t)
	url := fmt.Sprintf("http://%s", m.ClientListeners[0].Addr().String())
	l0 := createListener(t)
	l1 := createListener(t)
	addressMap := map[uint64]string{
		0: l0.Addr().String(),
		1: l1.Addr().String(),
	}

	pDataChan := make(chan *tDataBundle, 1)
	cDataChan := make(chan *tDataBundle, 1)
	// simulate two tasks on two nodes -- 0 and 1
	// 0 is parent, 1 is child
	f0 := &framework{
		name:     "framework_test_datarequest",
		etcdURLs: []string{url},
		taskID:   0,
		task: &testableTask{
			dataMap:  dataMap,
			dataChan: cDataChan,
		},
		topology:   NewTreeTopology(2, 1),
		ln:         l0,
		addressMap: addressMap,
	}
	f1 := &framework{
		name:     "framework_test_datarequest",
		etcdURLs: []string{url},
		taskID:   1,
		task: &testableTask{
			dataMap:  dataMap,
			dataChan: pDataChan,
		},
		topology:   NewTreeTopology(2, 1),
		ln:         l1,
		addressMap: addressMap,
	}
	f0.start()
	defer f0.stop()
	f1.start()
	defer f1.stop()

	for i, tt := range tests {
		// 0: F#DataRequest -> 1: T#ServeAsChild -> 0: T#ChildDataReady
		f0.DataRequest(1, tt.req)
		// from child(1)'s view at 1: T#ServeAsChild
		data := <-pDataChan
		expected := &tDataBundle{0, "", data.req, nil}
		if !reflect.DeepEqual(data, expected) {
			t.Errorf("#%d: data bundle want = %v, get = %v", i, expected, data)
		}
		// from parent(0)'s view at 0: T#ChildDataReady
		data = <-cDataChan
		expected = &tDataBundle{1, "", data.req, data.resp}
		if !reflect.DeepEqual(data, expected) {
			t.Errorf("#%d: data bundle want = %v, get = %v", i, expected, data)
		}

		// 1: F#DataRequest -> 0: T#ServeAsParent -> 1: T#ParentDataReady
		f1.DataRequest(0, tt.req)
		// from parent(0)'s view at 0: T#ServeAsParent
		data = <-cDataChan
		expected = &tDataBundle{1, "", data.req, nil}
		if !reflect.DeepEqual(data, expected) {
			t.Errorf("#%d: data bundle want = %v, get = %v", i, expected, data)
		}
		// from child(1)'s view at 1: T#ParentDataReady
		data = <-pDataChan
		expected = &tDataBundle{0, "", data.req, data.resp}
		if !reflect.DeepEqual(data, expected) {
			t.Errorf("#%d: data bundle want = %v, get = %v", i, expected, data)
		}
	}
}

type tDataBundle struct {
	id   uint64
	meta string
	req  string
	resp []byte
}

type testableTask struct {
	id        uint64
	framework Framework
	// dataMap will be used to serve data according to request
	dataMap map[string][]byte

	// This channel is used to convey data passed from framework back to the main
	// thread, for checking. Thus it's initialized and passed in from outside.
	//
	// The basic idea is that there are only two nodes -- one parent and one child.
	// When this channel is for parent, it passes information from child.
	dataChan chan *tDataBundle
}

func (t *testableTask) Init(taskID uint64, framework Framework, config Config) {
	t.id = taskID
	t.framework = framework
}
func (t *testableTask) Exit()                 {}
func (t *testableTask) SetEpoch(epoch uint64) {}

func (t *testableTask) ParentMetaReady(fromID uint64, meta string) {
	if t.dataChan != nil {
		t.dataChan <- &tDataBundle{fromID, meta, "", nil}
	}
}

func (t *testableTask) ChildMetaReady(fromID uint64, meta string) {
	t.ParentMetaReady(fromID, meta)
}

func (t *testableTask) ServeAsParent(fromID uint64, req string) []byte {
	if t.dataChan != nil {
		t.dataChan <- &tDataBundle{fromID, "", req, nil}
	}
	return t.dataMap[req]
}
func (t *testableTask) ServeAsChild(fromID uint64, req string) []byte {
	return t.ServeAsParent(fromID, req)
}
func (t *testableTask) ParentDataReady(fromID uint64, req string, resp []byte) {
	if t.dataChan != nil {
		t.dataChan <- &tDataBundle{fromID, "", req, resp}
	}
}

func (t *testableTask) ChildDataReady(fromID uint64, req string, resp []byte) {
	t.ParentDataReady(fromID, req, resp)
}

func createListener(t *testing.T) net.Listener {
	l, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.Listen(\"tcp4\", \"\") failed: %v", err)
	}
	return l
}
<|endoftext|>


<commit_before>package storage

import (
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	metav1beta1 "k8s.io/apimachinery/pkg/apis/meta/v1beta1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/duration"
	"k8s.io/kubernetes/pkg/printers"

	"github.com/operator-framework/operator-lifecycle-manager/pkg/package-server/apis/operators"
)

// translateTimestampSince returns the elapsed time since timestamp in
// human-readable approximation.
func translateTimestampSince(timestamp metav1.Time) string {
	if timestamp.IsZero() {
		return "<unknown>"
	}

	return duration.HumanDuration(time.Since(timestamp.Time))
}

func addTableHandlers(h printers.PrintHandler) {
	podColumnDefinitions := []metav1beta1.TableColumnDefinition{
		{Name: "Name", Type: "string", Format: "name", Description: metav1.ObjectMeta{}.SwaggerDoc()["name"]},
		{Name: "Catalog", Type: "string", Description: "The source catalog for this package"},
		{Name: "Age", Type: "string", Description: metav1.ObjectMeta{}.SwaggerDoc()["creationTimestamp"]},
	}
	h.TableHandler(podColumnDefinitions, printPackage)
	h.TableHandler(podColumnDefinitions, printPackageList)

}

func printPackage(manifest *operators.PackageManifest, options printers.PrintOptions) ([]metav1beta1.TableRow, error) {
	row := metav1beta1.TableRow{
		Object: runtime.RawExtension{Object: manifest},
	}
	row.Cells = append(row.Cells, manifest.Name, manifest.Status.CatalogSourceDisplayName, translateTimestampSince(manifest.CreationTimestamp))
	return []metav1beta1.TableRow{row}, nil
}

func printPackageList(manifestList *operators.PackageManifestList, options printers.PrintOptions) ([]metav1beta1.TableRow, error) {
	rows := make([]metav1beta1.TableRow, 0, len(manifestList.Items))
	for i := range manifestList.Items {
		r, err := printPackage(&manifestList.Items[i], options)
		if err != nil {
			return nil, err
		}
		rows = append(rows, r...)
	}
	return rows, nil
}
<commit_msg>Update PackageManfiest table handler<commit_after>package storage

import (
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	metav1beta1 "k8s.io/apimachinery/pkg/apis/meta/v1beta1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/duration"
	"k8s.io/kubernetes/pkg/printers"

	"github.com/operator-framework/operator-lifecycle-manager/pkg/package-server/apis/operators"
)

// translateTimestampSince returns the elapsed time since timestamp in
// human-readable approximation.
func translateTimestampSince(timestamp metav1.Time) string {
	if timestamp.IsZero() {
		return "<unknown>"
	}

	return duration.HumanDuration(time.Since(timestamp.Time))
}

func addTableHandlers(h printers.PrintHandler) {
	podColumnDefinitions := []metav1beta1.TableColumnDefinition{
		{Name: "Name", Type: "string", Format: "name", Description: metav1.ObjectMeta{}.SwaggerDoc()["name"]},
		{Name: "Catalog", Type: "string", Description: "The source catalog for this package"},
		{Name: "Age", Type: "string", Description: metav1.ObjectMeta{}.SwaggerDoc()["creationTimestamp"]},
	}
	h.TableHandler(podColumnDefinitions, printPackage)
	h.TableHandler(podColumnDefinitions, printPackageList)

}

func printPackage(manifest *operators.PackageManifest, options printers.GenerateOptions) ([]metav1beta1.TableRow, error) {
	row := metav1beta1.TableRow{
		Object: runtime.RawExtension{Object: manifest},
	}
	row.Cells = append(row.Cells, manifest.Name, manifest.Status.CatalogSourceDisplayName, translateTimestampSince(manifest.CreationTimestamp))
	return []metav1beta1.TableRow{row}, nil
}

func printPackageList(manifestList *operators.PackageManifestList, options printers.GenerateOptions) ([]metav1beta1.TableRow, error) {
	rows := make([]metav1beta1.TableRow, 0, len(manifestList.Items))
	for i := range manifestList.Items {
		r, err := printPackage(&manifestList.Items[i], options)
		if err != nil {
			return nil, err
		}
		rows = append(rows, r...)
	}
	return rows, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package route

import (
	"testing"
	"time"

	"golang.org/x/sync/errgroup"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/cache"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/ptr"
	"knative.dev/pkg/system"
	netv1alpha1 "knative.dev/serving/pkg/apis/networking/v1alpha1"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	fakeservingclient "knative.dev/serving/pkg/client/injection/client/fake"
	"knative.dev/serving/pkg/gc"
	"knative.dev/serving/pkg/network"
	"knative.dev/serving/pkg/reconciler/route/config"

	. "knative.dev/pkg/reconciler/testing"
	. "knative.dev/serving/pkg/testing/v1"
)

func TestNewRouteCallsSyncHandler(t *testing.T) {
	ctx, cancel, informers := SetupFakeContextWithCancel(t)

	// A standalone revision
	rev := getTestRevision("test-rev")
	// A route targeting the revision
	route := getTestRouteWithTrafficTargets(WithSpecTraffic(
		v1.TrafficTarget{
			RevisionName: "test-rev",
			Percent:      ptr.Int64(100),
		}))

	// Create fake clients
	configMapWatcher := configmap.NewStaticWatcher(&corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      config.DomainConfigName,
			Namespace: system.Namespace(),
		},
		Data: map[string]string{
			defaultDomainSuffix: "",
			prodDomainSuffix:    "selector:\n  app: prod",
		},
	}, &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      network.ConfigName,
			Namespace: system.Namespace(),
		},
		Data: map[string]string{},
	}, &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      gc.ConfigName,
			Namespace: system.Namespace(),
		},
		Data: map[string]string{},
	})

	ctrl := NewController(ctx, configMapWatcher)

	servingClient := fakeservingclient.Get(ctx)

	h := NewHooks()

	// Check for Ingress created as a signal that syncHandler ran
	h.OnCreate(&servingClient.Fake, "ingresses", func(obj runtime.Object) HookResult {
		ci := obj.(*netv1alpha1.Ingress)
		t.Logf("ingress created: %q", ci.Name)

		return HookComplete
	})

	eg := errgroup.Group{}

	waitInformers, err := controller.RunInformers(ctx.Done(), informers...)
	if err != nil {
		t.Fatalf("Failed to start informers: %v", err)
	}
	defer func() {
		cancel()
		if err := eg.Wait(); err != nil {
			t.Fatalf("Error running controller: %v", err)
		}
		waitInformers()
	}()

	// Run the controller.
	eg.Go(func() error {
		return ctrl.Run(2, ctx.Done())
	})

	if _, err := servingClient.ServingV1().Revisions(rev.Namespace).Create(rev); err != nil {
		t.Errorf("Unexpected error creating revision: %v", err)
	}

	for i, informer := range informers {
		if ok := cache.WaitForCacheSync(ctx.Done(), informer.HasSynced); !ok {
			t.Fatalf("failed to wait for cache at index %d to sync", i)
		}
	}

	if _, err := servingClient.ServingV1().Routes(route.Namespace).Create(route); err != nil {
		t.Errorf("Unexpected error creating route: %v", err)
	}

	if err := h.WaitForHooks(time.Second * 3); err != nil {
		t.Error(err)
	}
}
<commit_msg>Remove redundant cache sync and make route sync test more robust. (#7427)<commit_after>/*
Copyright 2018 The Knative Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package route

import (
	"testing"
	"time"

	"golang.org/x/sync/errgroup"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/ptr"
	"knative.dev/pkg/system"
	netv1alpha1 "knative.dev/serving/pkg/apis/networking/v1alpha1"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	fakeservingclient "knative.dev/serving/pkg/client/injection/client/fake"
	fakerevisioninformer "knative.dev/serving/pkg/client/injection/informers/serving/v1/revision/fake"
	fakerouteinformer "knative.dev/serving/pkg/client/injection/informers/serving/v1/route/fake"
	"knative.dev/serving/pkg/gc"
	"knative.dev/serving/pkg/network"
	"knative.dev/serving/pkg/reconciler/route/config"

	. "knative.dev/pkg/reconciler/testing"
	. "knative.dev/serving/pkg/testing/v1"
)

func TestNewRouteCallsSyncHandler(t *testing.T) {
	ctx, cancel, informers := SetupFakeContextWithCancel(t)

	// A standalone revision
	rev := getTestRevision("test-rev")
	// A route targeting the revision
	route := getTestRouteWithTrafficTargets(WithSpecTraffic(
		v1.TrafficTarget{
			RevisionName: "test-rev",
			Percent:      ptr.Int64(100),
		}))

	// Create fake clients
	configMapWatcher := configmap.NewStaticWatcher(&corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      config.DomainConfigName,
			Namespace: system.Namespace(),
		},
		Data: map[string]string{
			defaultDomainSuffix: "",
			prodDomainSuffix:    "selector:\n  app: prod",
		},
	}, &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      network.ConfigName,
			Namespace: system.Namespace(),
		},
		Data: map[string]string{},
	}, &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      gc.ConfigName,
			Namespace: system.Namespace(),
		},
		Data: map[string]string{},
	})

	servingClient := fakeservingclient.Get(ctx)

	h := NewHooks()
	// Check for Ingress created as a signal that syncHandler ran
	h.OnCreate(&servingClient.Fake, "ingresses", func(obj runtime.Object) HookResult {
		ci := obj.(*netv1alpha1.Ingress)
		t.Logf("ingress created: %q", ci.Name)

		return HookComplete
	})

	waitInformers, err := controller.RunInformers(ctx.Done(), informers...)
	if err != nil {
		t.Fatalf("Failed to start informers: %v", err)
	}

	// Run the controller.
	eg := errgroup.Group{}
	eg.Go(func() error {
		ctrl := NewController(ctx, configMapWatcher)
		return ctrl.Run(2, ctx.Done())
	})

	defer func() {
		cancel()
		if err := eg.Wait(); err != nil {
			t.Fatalf("Error running controller: %v", err)
		}
		waitInformers()
	}()

	if _, err := servingClient.ServingV1().Revisions(rev.Namespace).Create(rev); err != nil {
		t.Fatalf("Unexpected error creating revision: %v", err)
	}
	fakerevisioninformer.Get(ctx).Informer().GetIndexer().Add(rev)

	if _, err := servingClient.ServingV1().Routes(route.Namespace).Create(route); err != nil {
		t.Fatalf("Unexpected error creating route: %v", err)
	}
	fakerouteinformer.Get(ctx).Informer().GetIndexer().Add(route)

	if err := h.WaitForHooks(time.Second * 3); err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>