<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/jen20/awspolicyequivalence"
)

func TestAccAWSKmsKey_basic(t *testing.T) {
	var keyBefore, keyAfter kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &keyBefore),
				),
			},
			{
				Config: testAccAWSKmsKey_removedPolicy(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &keyAfter),
				),
			},
		},
	})
}

func TestAccAWSKmsKey_disappears(t *testing.T) {
	var key kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &key),
				),
			},
			{
				Config:             testAccAWSKmsKey_other_region(rName),
				PlanOnly:           true,
				ExpectNonEmptyPlan: true,
			},
		},
	})
}

func TestAccAWSKmsKey_policy(t *testing.T) {
	var key kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)
	expectedPolicyText := `{"Version":"2012-10-17","Id":"kms-tf-1","Statement":[{"Sid":"Enable IAM User Permissions","Effect":"Allow","Principal":{"AWS":"*"},"Action":"kms:*","Resource":"*"}]}`

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &key),
					testAccCheckAWSKmsKeyHasPolicy("aws_kms_key.foo", expectedPolicyText),
				),
			},
		},
	})
}

func TestAccAWSKmsKey_isEnabled(t *testing.T) {
	var key1, key2, key3 kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey_enabledRotation(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.bar", &key1),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "is_enabled", "true"),
					testAccCheckAWSKmsKeyIsEnabled(&key1, true),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "enable_key_rotation", "true"),
				),
			},
			{
				Config: testAccAWSKmsKey_disabled(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.bar", &key2),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "is_enabled", "false"),
					testAccCheckAWSKmsKeyIsEnabled(&key2, false),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "enable_key_rotation", "false"),
				),
			},
			{
				Config: testAccAWSKmsKey_enabled(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.bar", &key3),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "is_enabled", "true"),
					testAccCheckAWSKmsKeyIsEnabled(&key3, true),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "enable_key_rotation", "true"),
				),
			},
		},
	})
}

func TestAccAWSKmsKey_tags(t *testing.T) {
	var keyBefore kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey_tags(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &keyBefore),
					resource.TestCheckResourceAttr("aws_kms_key.foo", "tags.%", "2"),
				),
			},
		},
	})
}

func testAccCheckAWSKmsKeyHasPolicy(name string, expectedPolicyText string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No KMS Key ID is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).kmsconn

		out, err := conn.GetKeyPolicy(&kms.GetKeyPolicyInput{
			KeyId:      aws.String(rs.Primary.ID),
			PolicyName: aws.String("default"),
		})
		if err != nil {
			return err
		}

		actualPolicyText := *out.Policy

		equivalent, err := awspolicy.PoliciesAreEquivalent(actualPolicyText, expectedPolicyText)
		if err != nil {
			return fmt.Errorf("Error testing policy equivalence: %s", err)
		}
		if !equivalent {
			return fmt.Errorf("Non-equivalent policy error:\n\nexpected: %s\n\n     got: %s\n",
				expectedPolicyText, actualPolicyText)
		}

		return nil
	}
}

func testAccCheckAWSKmsKeyDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).kmsconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_kms_key" {
			continue
		}

		out, err := conn.DescribeKey(&kms.DescribeKeyInput{
			KeyId: aws.String(rs.Primary.ID),
		})

		if err != nil {
			return err
		}

		if *out.KeyMetadata.KeyState == "PendingDeletion" {
			return nil
		}

		return fmt.Errorf("KMS key still exists:\n%#v", out.KeyMetadata)
	}

	return nil
}

func testAccCheckAWSKmsKeyExists(name string, key *kms.KeyMetadata) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No KMS Key ID is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).kmsconn

		o, err := retryOnAwsCode("NotFoundException", func() (interface{}, error) {
			return conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(rs.Primary.ID),
			})
		})
		if err != nil {
			return err
		}
		out := o.(*kms.DescribeKeyOutput)

		*key = *out.KeyMetadata

		return nil
	}
}

func testAccCheckAWSKmsKeyIsEnabled(key *kms.KeyMetadata, isEnabled bool) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if *key.Enabled != isEnabled {
			return fmt.Errorf("Expected key %q to have is_enabled=%t, given %t",
				*key.Arn, isEnabled, *key.Enabled)
		}

		return nil
	}
}

func testAccAWSKmsKey(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "foo" {
    description = "Terraform acc test %s"
    deletion_window_in_days = 7
    policy = <<POLICY
{
  "Version": "2012-10-17",
  "Id": "kms-tf-1",
  "Statement": [
    {
      "Sid": "Enable IAM User Permissions",
      "Effect": "Allow",
      "Principal": {
        "AWS": "*"
      },
      "Action": "kms:*",
      "Resource": "*"
    }
  ]
}
POLICY
	tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_other_region(rName string) string {
	return fmt.Sprintf(`
provider "aws" {
	region = "us-east-1"
}
resource "aws_kms_key" "foo" {
    description = "Terraform acc test %s"
    deletion_window_in_days = 7
    policy = <<POLICY
{
  "Version": "2012-10-17",
  "Id": "kms-tf-1",
  "Statement": [
    {
      "Sid": "Enable IAM User Permissions",
      "Effect": "Allow",
      "Principal": {
        "AWS": "*"
      },
      "Action": "kms:*",
      "Resource": "*"
    }
  ]
}
POLICY
	tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_removedPolicy(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "foo" {
    description = "Terraform acc test %s"
    deletion_window_in_days = 7
    tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_enabledRotation(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "bar" {
    description = "Terraform acc test is_enabled %s"
    deletion_window_in_days = 7
    enable_key_rotation = true
    tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_disabled(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "bar" {
    description = "Terraform acc test is_enabled %s"
    deletion_window_in_days = 7
    enable_key_rotation = false
    is_enabled = false
    tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_enabled(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "bar" {
    description = "Terraform acc test is_enabled %s"
    deletion_window_in_days = 7
    enable_key_rotation = true
    is_enabled = true
    tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_tags(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "foo" {
    description = "Terraform acc test %s"
	tags {
		Name = "tf-acc-test-kms-key-%s"
		Key1 = "Value One"
		Description = "Very interesting"
	}
}`, rName, rName)
}
<commit_msg>Add sweepers for KMS keys<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/jen20/awspolicyequivalence"
)

func init() {
	resource.AddTestSweepers("aws_kms_key", &resource.Sweeper{
		Name: "aws_kms_key",
		F:    testSweepKmsKeys,
	})
}

func testSweepKmsKeys(region string) error {
	client, err := sharedClientForRegion(region)
	if err != nil {
		return fmt.Errorf("error getting client: %s", err)
	}
	conn := client.(*AWSClient).kmsconn

	err = conn.ListKeysPages(&kms.ListKeysInput{Limit: aws.Int64(int64(1000))}, func(out *kms.ListKeysOutput, lastPage bool) bool {
		for _, k := range out.Keys {
			kOut, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: k.KeyId,
			})
			if err != nil {
				log.Printf("Error: Failed to describe key %q: %s", *k.KeyId, err)
				return false
			}
			if *kOut.KeyMetadata.KeyManager == kms.KeyManagerTypeAws {
				// Skip (default) keys which are managed by AWS
				continue
			}
			if *kOut.KeyMetadata.KeyState == kms.KeyStatePendingDeletion {
				// Skip keys which are already scheduled for deletion
				continue
			}

			tOut, err := conn.ListResourceTags(&kms.ListResourceTagsInput{
				KeyId: k.KeyId,
			})
			if err != nil {
				log.Printf("Error: Failed to get tags for key %q: %s", *k.KeyId, err)
				return false
			}
			if !kmsTagHasPrefix(tOut.Tags, "Name", "tf-acc-test-kms-key-") {
				// Skip keys which don't have designated tag
				continue
			}

			_, err = conn.ScheduleKeyDeletion(&kms.ScheduleKeyDeletionInput{
				KeyId:               k.KeyId,
				PendingWindowInDays: aws.Int64(int64(7)),
			})
			if err != nil {
				log.Printf("Error: Failed to schedule key %q for deletion: %s", *k.KeyId, err)
				return false
			}
		}
		return !lastPage
	})
	if err != nil {
		return fmt.Errorf("Error describing KMS keys: %s", err)
	}

	return nil
}

func kmsTagHasPrefix(tags []*kms.Tag, key, prefix string) bool {
	for _, t := range tags {
		if *t.TagKey == key && strings.HasPrefix(*t.TagValue, prefix) {
			return true
		}
	}
	return false
}

func TestAccAWSKmsKey_basic(t *testing.T) {
	var keyBefore, keyAfter kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &keyBefore),
				),
			},
			{
				Config: testAccAWSKmsKey_removedPolicy(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &keyAfter),
				),
			},
		},
	})
}

func TestAccAWSKmsKey_disappears(t *testing.T) {
	var key kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &key),
				),
			},
			{
				Config:             testAccAWSKmsKey_other_region(rName),
				PlanOnly:           true,
				ExpectNonEmptyPlan: true,
			},
		},
	})
}

func TestAccAWSKmsKey_policy(t *testing.T) {
	var key kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)
	expectedPolicyText := `{"Version":"2012-10-17","Id":"kms-tf-1","Statement":[{"Sid":"Enable IAM User Permissions","Effect":"Allow","Principal":{"AWS":"*"},"Action":"kms:*","Resource":"*"}]}`

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &key),
					testAccCheckAWSKmsKeyHasPolicy("aws_kms_key.foo", expectedPolicyText),
				),
			},
		},
	})
}

func TestAccAWSKmsKey_isEnabled(t *testing.T) {
	var key1, key2, key3 kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey_enabledRotation(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.bar", &key1),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "is_enabled", "true"),
					testAccCheckAWSKmsKeyIsEnabled(&key1, true),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "enable_key_rotation", "true"),
				),
			},
			{
				Config: testAccAWSKmsKey_disabled(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.bar", &key2),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "is_enabled", "false"),
					testAccCheckAWSKmsKeyIsEnabled(&key2, false),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "enable_key_rotation", "false"),
				),
			},
			{
				Config: testAccAWSKmsKey_enabled(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.bar", &key3),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "is_enabled", "true"),
					testAccCheckAWSKmsKeyIsEnabled(&key3, true),
					resource.TestCheckResourceAttr("aws_kms_key.bar", "enable_key_rotation", "true"),
				),
			},
		},
	})
}

func TestAccAWSKmsKey_tags(t *testing.T) {
	var keyBefore kms.KeyMetadata
	rName := acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSKmsKeyDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSKmsKey_tags(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSKmsKeyExists("aws_kms_key.foo", &keyBefore),
					resource.TestCheckResourceAttr("aws_kms_key.foo", "tags.%", "3"),
				),
			},
		},
	})
}

func testAccCheckAWSKmsKeyHasPolicy(name string, expectedPolicyText string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No KMS Key ID is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).kmsconn

		out, err := conn.GetKeyPolicy(&kms.GetKeyPolicyInput{
			KeyId:      aws.String(rs.Primary.ID),
			PolicyName: aws.String("default"),
		})
		if err != nil {
			return err
		}

		actualPolicyText := *out.Policy

		equivalent, err := awspolicy.PoliciesAreEquivalent(actualPolicyText, expectedPolicyText)
		if err != nil {
			return fmt.Errorf("Error testing policy equivalence: %s", err)
		}
		if !equivalent {
			return fmt.Errorf("Non-equivalent policy error:\n\nexpected: %s\n\n     got: %s\n",
				expectedPolicyText, actualPolicyText)
		}

		return nil
	}
}

func testAccCheckAWSKmsKeyDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).kmsconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_kms_key" {
			continue
		}

		out, err := conn.DescribeKey(&kms.DescribeKeyInput{
			KeyId: aws.String(rs.Primary.ID),
		})

		if err != nil {
			return err
		}

		if *out.KeyMetadata.KeyState == "PendingDeletion" {
			return nil
		}

		return fmt.Errorf("KMS key still exists:\n%#v", out.KeyMetadata)
	}

	return nil
}

func testAccCheckAWSKmsKeyExists(name string, key *kms.KeyMetadata) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No KMS Key ID is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).kmsconn

		o, err := retryOnAwsCode("NotFoundException", func() (interface{}, error) {
			return conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(rs.Primary.ID),
			})
		})
		if err != nil {
			return err
		}
		out := o.(*kms.DescribeKeyOutput)

		*key = *out.KeyMetadata

		return nil
	}
}

func testAccCheckAWSKmsKeyIsEnabled(key *kms.KeyMetadata, isEnabled bool) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if *key.Enabled != isEnabled {
			return fmt.Errorf("Expected key %q to have is_enabled=%t, given %t",
				*key.Arn, isEnabled, *key.Enabled)
		}

		return nil
	}
}

func testAccAWSKmsKey(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "foo" {
    description = "Terraform acc test %s"
    deletion_window_in_days = 7
    policy = <<POLICY
{
  "Version": "2012-10-17",
  "Id": "kms-tf-1",
  "Statement": [
    {
      "Sid": "Enable IAM User Permissions",
      "Effect": "Allow",
      "Principal": {
        "AWS": "*"
      },
      "Action": "kms:*",
      "Resource": "*"
    }
  ]
}
POLICY
	tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_other_region(rName string) string {
	return fmt.Sprintf(`
provider "aws" {
	region = "us-east-1"
}
resource "aws_kms_key" "foo" {
    description = "Terraform acc test %s"
    deletion_window_in_days = 7
    policy = <<POLICY
{
  "Version": "2012-10-17",
  "Id": "kms-tf-1",
  "Statement": [
    {
      "Sid": "Enable IAM User Permissions",
      "Effect": "Allow",
      "Principal": {
        "AWS": "*"
      },
      "Action": "kms:*",
      "Resource": "*"
    }
  ]
}
POLICY
	tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_removedPolicy(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "foo" {
    description = "Terraform acc test %s"
    deletion_window_in_days = 7
    tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_enabledRotation(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "bar" {
    description = "Terraform acc test is_enabled %s"
    deletion_window_in_days = 7
    enable_key_rotation = true
    tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_disabled(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "bar" {
    description = "Terraform acc test is_enabled %s"
    deletion_window_in_days = 7
    enable_key_rotation = false
    is_enabled = false
    tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_enabled(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "bar" {
    description = "Terraform acc test is_enabled %s"
    deletion_window_in_days = 7
    enable_key_rotation = true
    is_enabled = true
    tags {
		Name = "tf-acc-test-kms-key-%s"
	}
}`, rName, rName)
}

func testAccAWSKmsKey_tags(rName string) string {
	return fmt.Sprintf(`
resource "aws_kms_key" "foo" {
    description = "Terraform acc test %s"
	tags {
		Name = "tf-acc-test-kms-key-%s"
		Key1 = "Value One"
		Description = "Very interesting"
	}
}`, rName, rName)
}
<|endoftext|>


<commit_before>package controllers

import (
	"encoding/json"
	"fmt"
	"github.com/orc/db"
	"github.com/orc/sessions"
	"github.com/orc/utils"
	"html/template"
	"reflect"
	"strconv"
	"time"
)

func (this *Handler) GetHistoryRequest() {
	if flag := sessions.CheackSession(this.Response, this.Request); !flag {
		return
	}
	this.Response.Header().Set("Access-Control-Allow-Origin", "*")
	this.Response.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
	this.Response.Header().Set("Content-type", "application/json")

	var data map[string]string
	decoder := json.NewDecoder(this.Request.Body)
	err := decoder.Decode(&data)
	utils.HandleErr("[Handler] Decode :", err, this.Response)

	event_id := data["event_id"]
	id := sessions.GetValue("id", this.Request).(string)

	users := GetModel("users")
	person, _ := users.Select([]string{"id", id}, "", []string{"person_id"})
	person_id := int(person[0].(map[string]interface{})["person_id"].(int64))

	query := `select param_id, p.name param_name, p.type, value, form_id, forms.name form_name from param_values
		inner join params p on param_values.param_id = p.id
		inner join forms on forms.id = p.form_id
		where person_id = $1 and event_id = $2;`

	rows := db.Query(query, []interface{}{person_id, event_id})
	rowsInf := db.Exec(query, []interface{}{person_id, event_id})

	size, _ := rowsInf.RowsAffected()
	columns, _ := rows.Columns()
	result := db.ConvertData(columns, size, rows)

	response, err := json.Marshal(result)
	utils.HandleErr("[Handle select] json.Marshal: ", err, nil)
	fmt.Fprintf(this.Response, "%s", string(response))
}

func (this *Handler) GetListHistoryEvents() {
	if flag := sessions.CheackSession(this.Response, this.Request); !flag {
		return
	}
	this.Response.Header().Set("Access-Control-Allow-Origin", "*")
	this.Response.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
	this.Response.Header().Set("Content-type", "application/json")

	var data map[string]interface{}
	decoder := json.NewDecoder(this.Request.Body)
	err := decoder.Decode(&data)
	utils.HandleErr("[Handler] Decode :", err, this.Response)

	id := sessions.GetValue("id", this.Request).(string)
	ids := utils.ArrayInterfaceToString(data["form_ids"].([]interface{}))

	users := GetModel("users")
	person, _ := users.Select([]string{"id", id}, "", []string{"person_id"})
	person_id := int(person[0].(map[string]interface{})["person_id"].(int64))

	model := GetModel("forms_types")
	result, _ := model.Select(ids, "OR", []string{"type_id"})
	//fmt.Println("result: ", result)

	query := `SELECT DISTINCT event_id, name FROM param_values
	inner join events on events.id = param_values.event_id
	WHERE event_id IN (SELECT DISTINCT event_id FROM events_types WHERE `

	var i int
	var params []interface{}

	for i = 1; i < reflect.ValueOf(result).Len(); i++ {
		query += "type_id=$" + strconv.Itoa(i) + " OR "
		params = append(params, result[i-1].(map[string]interface{})["type_id"])
	}

	query += "type_id=$" + strconv.Itoa(i) + ") AND person_id=$" + strconv.Itoa(i+1)

	params = append(params, result[i-1].(map[string]interface{})["type_id"])
	params = append(params, person_id)
	//fmt.Println("params: ", params)

	rows := db.Query(query, params)
	rowsInf := db.Exec(query, params)
	size, _ := rowsInf.RowsAffected()
	columns, _ := rows.Columns()
	events := db.ConvertData(columns, size, rows)

	response, err := json.Marshal(events)
	utils.HandleErr("[Handle GetListHistoryEvents] json.Marshal: ", err, nil)
	fmt.Fprintf(this.Response, "%s", string(response))
}

func (this *Handler) SaveUserRequest() {
	if flag := sessions.CheackSession(this.Response, this.Request); !flag {
		return
	}
	this.Response.Header().Set("Access-Control-Allow-Origin", "*")
	this.Response.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
	this.Response.Header().Set("Content-type", "application/json")

	var data map[string]interface{}
	decoder := json.NewDecoder(this.Request.Body)
	err := decoder.Decode(&data)
	utils.HandleErr("[Handler] Decode :", err, this.Response)

	event_id := int(data["event_id"].(float64))
	id := sessions.GetValue("id", this.Request).(string)

	users := GetModel("users")
	person, _ := users.Select([]string{"id", id}, "", []string{"person_id"})
	person_id := int(person[0].(map[string]interface{})["person_id"].(int64))

	persons_events := GetModel("persons_events")
	person, _ = persons_events.Select([]string{"person_id", strconv.Itoa(person_id), "event_id", strconv.Itoa(event_id)}, "AND", []string{"person_id"})

	var response interface{}
	inf := data["data"].([]interface{})
	param_values := GetModel("param_values")
	t := time.Now()

	if len(person) == 0 {
		persons_events.Insert(
			[]string{"person_id", "event_id", "reg_date", "last_date"},
			[]interface{}{person_id, event_id,
				t.Format("2006-01-02"),
				t.Format("2006-01-02")})

		for _, element := range inf {
			param_id := element.(map[string]interface{})["name"]
			value := element.(map[string]interface{})["value"]
			param_values.Insert(
				[]string{"person_id", "event_id", "param_id", "value"},
				[]interface{}{person_id, event_id, param_id, value})
		}
		response = map[string]interface{}{"result": "ok"}
	} else if len(person) != 0 {
		for _, element := range inf {
			param_id := element.(map[string]interface{})["name"]
			value := element.(map[string]interface{})["value"]
			param_values.Update(
				[]string{"value"},
				[]interface{}{value, person_id, event_id, param_id},
				"person_id=$"+strconv.Itoa(2)+" AND event_id=$"+strconv.Itoa(3)+" AND param_id=$"+strconv.Itoa(4))
		}
		persons_events.Update(
			[]string{"last_date"},
			[]interface{}{t.Format("2006-01-02"),
				person_id, event_id},
			"person_id=$"+strconv.Itoa(2)+" AND event_id=$"+strconv.Itoa(3))
		response = map[string]interface{}{"result": "ok"}
	} else {
		response = map[string]interface{}{"result": "exists"}
	}

	result, err := json.Marshal(response)
	utils.HandleErr("[Handle select] json.Marshal: ", err, nil)
	fmt.Fprintf(this.Response, "%s", string(result))
}

func (this *Handler) GetRequest(tableName, id string) {
	tmp, err := template.ParseFiles(
		"mvc/views/item.html",
		"mvc/views/header.html",
		"mvc/views/footer.html")
	utils.HandleErr("[Handler.Show] template.ParseFiles: ", err, nil)

	reaponse, err := json.Marshal(MegoJoin(tableName, id))
	utils.HandleErr("[Handler.Show] template.json.Marshal: ", err, nil)

	err = tmp.ExecuteTemplate(this.Response, "item", template.JS(reaponse))
	utils.HandleErr("[Handler.Show] tmp.Execute: ", err, nil)
}

func MegoJoin(tableName, id string) RequestModel {
	var E []interface{}
	var T []interface{}
	var F []interface{}
	var P []interface{}

	E = db.Select("events", []string{"id", id}, "", []string{"id", "name"})

	query := db.InnerJoin(
		[]string{"id", "name"},
		"t",
		"events_types",
		"e_t",
		[]string{"event_id", "type_id"},
		[]string{"events", "event_types"},
		[]string{"e", "t"},
		[]string{"id", "id"},
		"where e.id=$1")

	rows := db.Query(query, []interface{}{id})
	rowsInf := db.Exec(query, []interface{}{id})
	l, _ := rowsInf.RowsAffected()
	c, _ := rows.Columns()
	T = db.ConvertData(c, l, rows)

	for i := 0; i < len(T); i++ {
		id := T[i].(map[string]interface{})["id"]

		query := db.InnerJoin(
			[]string{"id", "name"},
			"f",
			"forms_types",
			"f_t",
			[]string{"form_id", "type_id"},
			[]string{"forms", "event_types"},
			[]string{"f", "t"},
			[]string{"id", "id"},
			"where t.id=$1")

		rows := db.Query(query, []interface{}{id})
		rowsInf := db.Exec(query, []interface{}{id})
		l, _ := rowsInf.RowsAffected()
		c, _ := rows.Columns()
		F = append(F, db.ConvertData(c, l, rows))
	}

	for i := 0; i < len(F); i++ {
		var PP []interface{}
		for j := 0; j < len(F[i].([]interface{})); j++ {
			item := F[i].([]interface{})[j]
			id := item.(map[string]interface{})["id"]

			query := db.InnerJoin(
				[]string{"id", "name", "type"},
				"p",
				"params",
				"p",
				[]string{"form_id"},
				[]string{"forms"},
				[]string{"f"},
				[]string{"id"},
				"where f.id=$1")

			rows := db.Query(query, []interface{}{id})
			rowsInf := db.Exec(query, []interface{}{id})
			l, _ := rowsInf.RowsAffected()
			c, _ := rows.Columns()
			PP = append(PP, db.ConvertData(c, l, rows))
		}
		P = append(P, PP)
	}
	return RequestModel{E: E, T: T, F: F, P: P}
}
<commit_msg>controllers: fix whitespace<commit_after>package controllers

import (
    "encoding/json"
    "fmt"
    "github.com/orc/db"
    "github.com/orc/sessions"
    "github.com/orc/utils"
    "html/template"
    "reflect"
    "strconv"
    "time"
)

func (this *Handler) GetHistoryRequest() {
    if flag := sessions.CheackSession(this.Response, this.Request); !flag {
        return
    }
    this.Response.Header().Set("Access-Control-Allow-Origin", "*")
    this.Response.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
    this.Response.Header().Set("Content-type", "application/json")

    var data map[string]string
    decoder := json.NewDecoder(this.Request.Body)
    err := decoder.Decode(&data)
    utils.HandleErr("[Handler] Decode :", err, this.Response)

    event_id := data["event_id"]
    id := sessions.GetValue("id", this.Request).(string)

    users := GetModel("users")
    person, _ := users.Select([]string{"id", id}, "", []string{"person_id"})
    person_id := int(person[0].(map[string]interface{})["person_id"].(int64))

    query := `select param_id, p.name param_name, p.type, value, form_id, forms.name form_name from param_values
        inner join params p on param_values.param_id = p.id
        inner join forms on forms.id = p.form_id
        where person_id = $1 and event_id = $2;`

    rows := db.Query(query, []interface{}{person_id, event_id})
    rowsInf := db.Exec(query, []interface{}{person_id, event_id})

    size, _ := rowsInf.RowsAffected()
    columns, _ := rows.Columns()
    result := db.ConvertData(columns, size, rows)

    response, err := json.Marshal(result)
    utils.HandleErr("[Handle select] json.Marshal: ", err, nil)
    fmt.Fprintf(this.Response, "%s", string(response))
}

func (this *Handler) GetListHistoryEvents() {
    if flag := sessions.CheackSession(this.Response, this.Request); !flag {
        return
    }
    this.Response.Header().Set("Access-Control-Allow-Origin", "*")
    this.Response.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
    this.Response.Header().Set("Content-type", "application/json")

    var data map[string]interface{}
    decoder := json.NewDecoder(this.Request.Body)
    err := decoder.Decode(&data)
    utils.HandleErr("[Handler] Decode :", err, this.Response)

    id := sessions.GetValue("id", this.Request).(string)
    ids := utils.ArrayInterfaceToString(data["form_ids"].([]interface{}))

    users := GetModel("users")
    person, _ := users.Select([]string{"id", id}, "", []string{"person_id"})
    person_id := int(person[0].(map[string]interface{})["person_id"].(int64))

    model := GetModel("forms_types")
    result, _ := model.Select(ids, "OR", []string{"type_id"})
    //fmt.Println("result: ", result)

    query := `SELECT DISTINCT event_id, name FROM param_values
    inner join events on events.id = param_values.event_id
    WHERE event_id IN (SELECT DISTINCT event_id FROM events_types WHERE `

    var i int
    var params []interface{}

    for i = 1; i < reflect.ValueOf(result).Len(); i++ {
        query += "type_id=$" + strconv.Itoa(i) + " OR "
        params = append(params, result[i-1].(map[string]interface{})["type_id"])
    }

    query += "type_id=$" + strconv.Itoa(i) + ") AND person_id=$" + strconv.Itoa(i+1)

    params = append(params, result[i-1].(map[string]interface{})["type_id"])
    params = append(params, person_id)
    //fmt.Println("params: ", params)

    rows := db.Query(query, params)
    rowsInf := db.Exec(query, params)
    size, _ := rowsInf.RowsAffected()
    columns, _ := rows.Columns()
    events := db.ConvertData(columns, size, rows)

    response, err := json.Marshal(events)
    utils.HandleErr("[Handle GetListHistoryEvents] json.Marshal: ", err, nil)
    fmt.Fprintf(this.Response, "%s", string(response))
}

func (this *Handler) SaveUserRequest() {
    if flag := sessions.CheackSession(this.Response, this.Request); !flag {
        return
    }
    this.Response.Header().Set("Access-Control-Allow-Origin", "*")
    this.Response.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
    this.Response.Header().Set("Content-type", "application/json")

    var data map[string]interface{}
    decoder := json.NewDecoder(this.Request.Body)
    err := decoder.Decode(&data)
    utils.HandleErr("[Handler] Decode :", err, this.Response)

    event_id := int(data["event_id"].(float64))
    id := sessions.GetValue("id", this.Request).(string)

    users := GetModel("users")
    person, _ := users.Select([]string{"id", id}, "", []string{"person_id"})
    person_id := int(person[0].(map[string]interface{})["person_id"].(int64))

    persons_events := GetModel("persons_events")
    person, _ = persons_events.Select([]string{"person_id", strconv.Itoa(person_id), "event_id", strconv.Itoa(event_id)}, "AND", []string{"person_id"})

    var response interface{}
    inf := data["data"].([]interface{})
    param_values := GetModel("param_values")
    t := time.Now()

    if len(person) == 0 {
        persons_events.Insert(
            []string{"person_id", "event_id", "reg_date", "last_date"},
            []interface{}{person_id, event_id,
                t.Format("2006-01-02"),
                t.Format("2006-01-02")})

        for _, element := range inf {
            param_id := element.(map[string]interface{})["name"]
            value := element.(map[string]interface{})["value"]
            param_values.Insert(
                []string{"person_id", "event_id", "param_id", "value"},
                []interface{}{person_id, event_id, param_id, value})
        }
        response = map[string]interface{}{"result": "ok"}
    } else if len(person) != 0 {
        for _, element := range inf {
            param_id := element.(map[string]interface{})["name"]
            value := element.(map[string]interface{})["value"]
            param_values.Update(
                []string{"value"},
                []interface{}{value, person_id, event_id, param_id},
                "person_id=$"+strconv.Itoa(2)+" AND event_id=$"+strconv.Itoa(3)+" AND param_id=$"+strconv.Itoa(4))
        }
        persons_events.Update(
            []string{"last_date"},
            []interface{}{t.Format("2006-01-02"),
                person_id, event_id},
            "person_id=$"+strconv.Itoa(2)+" AND event_id=$"+strconv.Itoa(3))
        response = map[string]interface{}{"result": "ok"}
    } else {
        response = map[string]interface{}{"result": "exists"}
    }

    result, err := json.Marshal(response)
    utils.HandleErr("[Handle select] json.Marshal: ", err, nil)
    fmt.Fprintf(this.Response, "%s", string(result))
}

func (this *Handler) GetRequest(tableName, id string) {
    tmp, err := template.ParseFiles(
        "mvc/views/item.html",
        "mvc/views/header.html",
        "mvc/views/footer.html")
    utils.HandleErr("[Handler.Show] template.ParseFiles: ", err, nil)

    reaponse, err := json.Marshal(MegoJoin(tableName, id))
    utils.HandleErr("[Handler.Show] template.json.Marshal: ", err, nil)

    err = tmp.ExecuteTemplate(this.Response, "item", template.JS(reaponse))
    utils.HandleErr("[Handler.Show] tmp.Execute: ", err, nil)
}

func MegoJoin(tableName, id string) RequestModel {
    var E []interface{}
    var T []interface{}
    var F []interface{}
    var P []interface{}

    E = db.Select("events", []string{"id", id}, "", []string{"id", "name"})

    query := db.InnerJoin(
        []string{"id", "name"},
        "t",
        "events_types",
        "e_t",
        []string{"event_id", "type_id"},
        []string{"events", "event_types"},
        []string{"e", "t"},
        []string{"id", "id"},
        "where e.id=$1")

    rows := db.Query(query, []interface{}{id})
    rowsInf := db.Exec(query, []interface{}{id})
    l, _ := rowsInf.RowsAffected()
    c, _ := rows.Columns()
    T = db.ConvertData(c, l, rows)

    for i := 0; i < len(T); i++ {
        id := T[i].(map[string]interface{})["id"]

        query := db.InnerJoin(
            []string{"id", "name"},
            "f",
            "forms_types",
            "f_t",
            []string{"form_id", "type_id"},
            []string{"forms", "event_types"},
            []string{"f", "t"},
            []string{"id", "id"},
            "where t.id=$1")

        rows := db.Query(query, []interface{}{id})
        rowsInf := db.Exec(query, []interface{}{id})
        l, _ := rowsInf.RowsAffected()
        c, _ := rows.Columns()
        F = append(F, db.ConvertData(c, l, rows))
    }

    for i := 0; i < len(F); i++ {
        var PP []interface{}
        for j := 0; j < len(F[i].([]interface{})); j++ {
            item := F[i].([]interface{})[j]
            id := item.(map[string]interface{})["id"]

            query := db.InnerJoin(
                []string{"id", "name", "type"},
                "p",
                "params",
                "p",
                []string{"form_id"},
                []string{"forms"},
                []string{"f"},
                []string{"id"},
                "where f.id=$1")

            rows := db.Query(query, []interface{}{id})
            rowsInf := db.Exec(query, []interface{}{id})
            l, _ := rowsInf.RowsAffected()
            c, _ := rows.Columns()
            PP = append(PP, db.ConvertData(c, l, rows))
        }
        P = append(P, PP)
    }
    return RequestModel{E: E, T: T, F: F, P: P}
}
<|endoftext|>


<commit_before>// Package db manages the queuing/persistance for the postmaster
package db

import (
	"encoding/json"
	"errors"
	"time"

	"github.com/levenlabs/go-llog"
	"github.com/levenlabs/golib/genapi"
	"github.com/levenlabs/postmaster/ga"
	"github.com/levenlabs/postmaster/sender"
	"github.com/mediocregopher/okq-go.v2"
)

var (
	normalQueue     = "email-normal"
	statsQueue      = "stats-normal"
	uniqueArgStatID = "pmStatsID"
	uniqueArgEnvID  = "pmEnvID"
)

var jobCh chan job

var useOkq bool

type job struct {
	Queue    string
	Contents string
	RespCh   chan error
}

func init() {
	ga.GA.AppendInit(func(g *genapi.GenAPI) {
		if ga.GA.OkqInfo.Client == nil {
			return
		}
		jobCh = make(chan job)

		okqClient := ga.GA.OkqInfo.Client
		// Receive jobs from StoreSendJob() and StoreStatsJob() and Push into okq
		go func() {
			for job := range jobCh {
				job.RespCh <- okqClient.Push(job.Queue, job.Contents, okq.Normal)
			}
		}()

		// Receive jobs from okq and send to sender
		consumeSpin(handleSendEvent, normalQueue)

		// Receive jobs from okq and store in stats
		consumeSpin(handleStatsEvent, statsQueue)

		useOkq = true
	})
}

// DisableOkq turns off using okq for job storing
// this should ONLY be called during testing
func DisableOkq() {
	useOkq = false
}

func consumeSpin(fn okq.ConsumerFunc, q string) {
	llog.Info("creating okq consumer", llog.KV{"queue": q})
	consumer := ga.GA.OkqInfo.Client
	go func(c *okq.Client) {
		for {
			err := <-c.Consumer(fn, nil, q)
			llog.Error("consumer error", llog.KV{"queue": q}, llog.ErrKV(err))
			time.Sleep(10 * time.Second)
		}
	}(consumer)
}

// StoreSendJob creates a new Mail job with jobContents and sends it to okq
func StoreSendJob(jobContents string) error {
	if !useOkq {
		if !sendEmail(jobContents) {
			return errors.New("Failed to send email (bypassing okq)")
		}
		return nil
	}
	respCh := make(chan error)
	jobCh <- job{normalQueue, jobContents, respCh}
	return <-respCh
}

// StoreStatsJob creates a new statsJob with jobContents and sends it to okq
func StoreStatsJob(jobContents string) error {
	if !useOkq {
		if !storeStats(jobContents) {
			return errors.New("Failed to store stats (bypassing okq)")
		}
		return nil
	}
	respCh := make(chan error)
	jobCh <- job{statsQueue, jobContents, respCh}
	return <-respCh
}

func handleSendEvent(e okq.Event) bool {
	return sendEmail(e.Contents)
}

func handleStatsEvent(e okq.Event) bool {
	return storeStats(e.Contents)
}

func sendEmail(jobContents string) bool {
	job := new(sender.Mail)
	err := json.Unmarshal([]byte(jobContents), job)
	if err != nil {
		llog.Error("error json decoding into sender.Mail", llog.KV{
			"jobContents": jobContents,
		}, llog.ErrKV(err))
		// since we cannot process this job, no reason to have it keep around
		return true
	}

	env := ga.Environment
	id := GenerateEmailID(job.To, job.Flags, job.UniqueID, env)
	if id != "" {
		if job.UniqueArgs == nil {
			job.UniqueArgs = make(map[string]string)
		}
		job.UniqueArgs[uniqueArgStatID] = id
		job.UniqueArgs[uniqueArgEnvID] = env
	}

	llog.Info("processing send job", llog.KV{"id": id, "recipient": job.To})
	err = sender.Send(job)
	if err != nil {
		if id != "" {
			// if we ran into an error sending the email, delete the emailID
			rerr := removeEmailID(id)
			if rerr != nil {
				llog.Error("error deleting failed emailID", llog.KV{"id": id}, llog.ErrKV(err))
			}
		}

		llog.Error("error calling sender.Send", llog.KV{"jobContents": jobContents, "id": id}, llog.ErrKV(err))
		return false
	}
	return true
}

func logMarkError(err error, kv llog.KV) {
	if err != nil {
		llog.Error("error marking email", kv, llog.ErrKV(err))
	}
}

func storeStats(jobContents string) bool {
	job := new(StatsJob)
	err := json.Unmarshal([]byte(jobContents), job)
	if err != nil {
		llog.Error("error json decoding into StatsJob", llog.KV{
			"jobContents": jobContents,
			"err":         err,
		})
		// since we cannot process this job, no reason to have it keep around
		return true
	}

	kv := llog.KV{
		"id":     job.StatsID,
		"type":   job.Type,
		"reason": job.Reason,
		"email":  job.Email,
	}
	llog.Info("processing stats job", kv)
	switch job.Type {
	case "delivered":
		err = MarkAsDelivered(job.StatsID)
		logMarkError(err, kv)
	case "open":
		err = MarkAsOpened(job.StatsID)
		logMarkError(err, kv)
	case "bounce":
		err = MarkAsBounced(job.StatsID, job.Reason)
		logMarkError(err, kv)

		err = StoreEmailBounce(job.Email)
		if err != nil {
			llog.Error("error storing email as bounced", kv, llog.ErrKV(err))
		}
	case "spamreport":
		err = MarkAsSpamReported(job.StatsID)
		logMarkError(err, kv)

		err = StoreEmailSpam(job.Email)
		if err != nil {
			llog.Error("error storing email as spamed", kv, llog.ErrKV(err))
		}
	case "dropped":
		//depending on the reason we should mark the email as invalid
		err = MarkAsDropped(job.StatsID, job.Reason)
		logMarkError(err, kv)
	default:
		llog.Warn("received unknown job type", llog.KV{"type": job.Type})
	}
	return true
}
<commit_msg>fix some backwards incompatible changes to okq driver<commit_after>// Package db manages the queuing/persistance for the postmaster
package db

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	"github.com/levenlabs/go-llog"
	"github.com/levenlabs/golib/genapi"
	"github.com/levenlabs/postmaster/ga"
	"github.com/levenlabs/postmaster/sender"
	"github.com/mediocregopher/okq-go.v2"
)

var (
	normalQueue     = "email-normal"
	statsQueue      = "stats-normal"
	uniqueArgStatID = "pmStatsID"
	uniqueArgEnvID  = "pmEnvID"
)

var jobCh chan job

var useOkq bool

type job struct {
	Queue    string
	Contents string
	RespCh   chan error
}

func init() {
	ga.GA.AppendInit(func(g *genapi.GenAPI) {
		if ga.GA.OkqInfo.Client == nil {
			return
		}
		jobCh = make(chan job)

		okqClient := ga.GA.OkqInfo.Client
		// Receive jobs from StoreSendJob() and StoreStatsJob() and Push into okq
		go func() {
			for job := range jobCh {
				job.RespCh <- okqClient.Push(job.Queue, job.Contents, okq.Normal)
			}
		}()

		// Receive jobs from okq and send to sender
		consumeSpin(handleSendEvent, normalQueue)

		// Receive jobs from okq and store in stats
		consumeSpin(handleStatsEvent, statsQueue)

		useOkq = true
	})
}

// DisableOkq turns off using okq for job storing
// this should ONLY be called during testing
func DisableOkq() {
	useOkq = false
}

func consumeSpin(fn okq.ConsumerFunc, q string) {
	llog.Info("creating okq consumer", llog.KV{"queue": q})
	consumer := ga.GA.OkqInfo.Client
	go func(c *okq.Client) {
		for {
			err := <-c.Consumer(fn, nil, q)
			llog.Error("consumer error", llog.KV{"queue": q}, llog.ErrKV(err))
			time.Sleep(10 * time.Second)
		}
	}(consumer)
}

// StoreSendJob creates a new Mail job with jobContents and sends it to okq
func StoreSendJob(jobContents string) error {
	if !useOkq {
		if !sendEmail(jobContents) {
			return errors.New("Failed to send email (bypassing okq)")
		}
		return nil
	}
	respCh := make(chan error)
	jobCh <- job{normalQueue, jobContents, respCh}
	return <-respCh
}

// StoreStatsJob creates a new statsJob with jobContents and sends it to okq
func StoreStatsJob(jobContents string) error {
	if !useOkq {
		if !storeStats(jobContents) {
			return errors.New("Failed to store stats (bypassing okq)")
		}
		return nil
	}
	respCh := make(chan error)
	jobCh <- job{statsQueue, jobContents, respCh}
	return <-respCh
}

func handleSendEvent(_ context.Context, e okq.Event) bool {
	return sendEmail(e.Contents)
}

func handleStatsEvent(_ context.Context, e okq.Event) bool {
	return storeStats(e.Contents)
}

func sendEmail(jobContents string) bool {
	job := new(sender.Mail)
	err := json.Unmarshal([]byte(jobContents), job)
	if err != nil {
		llog.Error("error json decoding into sender.Mail", llog.KV{
			"jobContents": jobContents,
		}, llog.ErrKV(err))
		// since we cannot process this job, no reason to have it keep around
		return true
	}

	env := ga.Environment
	id := GenerateEmailID(job.To, job.Flags, job.UniqueID, env)
	if id != "" {
		if job.UniqueArgs == nil {
			job.UniqueArgs = make(map[string]string)
		}
		job.UniqueArgs[uniqueArgStatID] = id
		job.UniqueArgs[uniqueArgEnvID] = env
	}

	llog.Info("processing send job", llog.KV{"id": id, "recipient": job.To})
	err = sender.Send(job)
	if err != nil {
		if id != "" {
			// if we ran into an error sending the email, delete the emailID
			rerr := removeEmailID(id)
			if rerr != nil {
				llog.Error("error deleting failed emailID", llog.KV{"id": id}, llog.ErrKV(err))
			}
		}

		llog.Error("error calling sender.Send", llog.KV{"jobContents": jobContents, "id": id}, llog.ErrKV(err))
		return false
	}
	return true
}

func logMarkError(err error, kv llog.KV) {
	if err != nil {
		llog.Error("error marking email", kv, llog.ErrKV(err))
	}
}

func storeStats(jobContents string) bool {
	job := new(StatsJob)
	err := json.Unmarshal([]byte(jobContents), job)
	if err != nil {
		llog.Error("error json decoding into StatsJob", llog.KV{
			"jobContents": jobContents,
			"err":         err,
		})
		// since we cannot process this job, no reason to have it keep around
		return true
	}

	kv := llog.KV{
		"id":     job.StatsID,
		"type":   job.Type,
		"reason": job.Reason,
		"email":  job.Email,
	}
	llog.Info("processing stats job", kv)
	switch job.Type {
	case "delivered":
		err = MarkAsDelivered(job.StatsID)
		logMarkError(err, kv)
	case "open":
		err = MarkAsOpened(job.StatsID)
		logMarkError(err, kv)
	case "bounce":
		err = MarkAsBounced(job.StatsID, job.Reason)
		logMarkError(err, kv)

		err = StoreEmailBounce(job.Email)
		if err != nil {
			llog.Error("error storing email as bounced", kv, llog.ErrKV(err))
		}
	case "spamreport":
		err = MarkAsSpamReported(job.StatsID)
		logMarkError(err, kv)

		err = StoreEmailSpam(job.Email)
		if err != nil {
			llog.Error("error storing email as spamed", kv, llog.ErrKV(err))
		}
	case "dropped":
		//depending on the reason we should mark the email as invalid
		err = MarkAsDropped(job.StatsID, job.Reason)
		logMarkError(err, kv)
	default:
		llog.Warn("received unknown job type", llog.KV{"type": job.Type})
	}
	return true
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach 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.
//
// Author: Peter Mattis (peter@cockroachlabs.com)

package client_test

import (
	"fmt"
	"reflect"
	"strings"
	"testing"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/util/caller"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/log"
)

func setup() (server.TestServer, *client.DB) {
	s := server.StartTestServer(nil)
	return s, s.DB()
}

func ExampleDB_Get() {
	s, db := setup()
	defer s.Stop()

	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	fmt.Printf("aa=%s\n", result.ValueBytes())

	// Output:
	// aa=
}

func ExampleDB_Put() {
	s, db := setup()
	defer s.Stop()

	if err := db.Put("aa", "1"); err != nil {
		panic(err)
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	fmt.Printf("aa=%s\n", result.ValueBytes())

	// Output:
	// aa=1
}

func ExampleDB_CPut() {
	s, db := setup()
	defer s.Stop()

	if err := db.Put("aa", "1"); err != nil {
		panic(err)
	}
	if err := db.CPut("aa", "2", "1"); err != nil {
		panic(err)
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	fmt.Printf("aa=%s\n", result.ValueBytes())

	if err = db.CPut("aa", "3", "1"); err == nil {
		panic("expected error from conditional put")
	}
	result, err = db.Get("aa")
	if err != nil {
		panic(err)
	}
	fmt.Printf("aa=%s\n", result.ValueBytes())

	if err = db.CPut("bb", "4", "1"); err == nil {
		panic("expected error from conditional put")
	}
	result, err = db.Get("bb")
	if err != nil {
		panic(err)
	}
	fmt.Printf("bb=%s\n", result.ValueBytes())
	if err = db.CPut("bb", "4", nil); err != nil {
		panic(err)
	}
	result, err = db.Get("bb")
	if err != nil {
		panic(err)
	}
	fmt.Printf("bb=%s\n", result.ValueBytes())

	// Output:
	// aa=2
	// aa=2
	// bb=
	// bb=4
}

func ExampleDB_InitPut() {
	s, db := setup()
	defer s.Stop()

	if err := db.InitPut("aa", "1"); err != nil {
		panic(err)
	}
	if err := db.InitPut("aa", "1"); err != nil {
		panic(err)
	}
	if err := db.InitPut("aa", "2"); err == nil {
		panic("expected error from init put")
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	fmt.Printf("aa=%s\n", result.ValueBytes())

	// Output:
	// aa=1
}

func ExampleDB_Inc() {
	s, db := setup()
	defer s.Stop()

	if _, err := db.Inc("aa", 100); err != nil {
		panic(err)
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	fmt.Printf("aa=%d\n", result.ValueInt())

	// Output:
	// aa=100
}

func ExampleBatch() {
	s, db := setup()
	defer s.Stop()

	b := &client.Batch{}
	b.Get("aa")
	b.Put("bb", "2")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	for _, result := range b.Results {
		for _, row := range result.Rows {
			fmt.Printf("%s=%s\n", row.Key, row.ValueBytes())
		}
	}

	// Output:
	// "aa"=
	// "bb"=2
}

func ExampleDB_Scan() {
	s, db := setup()
	defer s.Stop()

	b := &client.Batch{}
	b.Put("aa", "1")
	b.Put("ab", "2")
	b.Put("bb", "3")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	rows, err := db.Scan("a", "b", 100)
	if err != nil {
		panic(err)
	}
	for i, row := range rows {
		fmt.Printf("%d: %s=%s\n", i, row.Key, row.ValueBytes())
	}

	// Output:
	// 0: "aa"=1
	// 1: "ab"=2
}

func ExampleDB_ReverseScan() {
	s, db := setup()
	defer s.Stop()

	b := &client.Batch{}
	b.Put("aa", "1")
	b.Put("ab", "2")
	b.Put("bb", "3")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	rows, err := db.ReverseScan("ab", "c", 100)
	if err != nil {
		panic(err)
	}
	for i, row := range rows {
		fmt.Printf("%d: %s=%s\n", i, row.Key, row.ValueBytes())
	}

	// Output:
	// 0: "bb"=3
	// 1: "ab"=2
}

func ExampleDB_Del() {
	s, db := setup()
	defer s.Stop()

	b := &client.Batch{}
	b.Put("aa", "1")
	b.Put("ab", "2")
	b.Put("ac", "3")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	if err := db.Del("ab"); err != nil {
		panic(err)
	}
	rows, err := db.Scan("a", "b", 100)
	if err != nil {
		panic(err)
	}
	for i, row := range rows {
		fmt.Printf("%d: %s=%s\n", i, row.Key, row.ValueBytes())
	}

	// Output:
	// 0: "aa"=1
	// 1: "ac"=3
}

func ExampleTxn_Commit() {
	s, db := setup()
	defer s.Stop()

	err := db.Txn(func(txn *client.Txn) error {
		b := txn.NewBatch()
		b.Put("aa", "1")
		b.Put("ab", "2")
		return txn.CommitInBatch(b)
	})
	if err != nil {
		panic(err)
	}

	b := &client.Batch{}
	b.Get("aa")
	b.Get("ab")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	for i, result := range b.Results {
		for j, row := range result.Rows {
			fmt.Printf("%d/%d: %s=%s\n", i, j, row.Key, row.ValueBytes())
		}
	}

	// Output:
	// 0/0: "aa"=1
	// 1/0: "ab"=2
}

func ExampleDB_Put_insecure() {
	ctx := server.MakeTestContext()
	ctx.Insecure = true
	s := server.TestServer{
		Ctx: &ctx,
	}
	if err := s.Start(); err != nil {
		log.Fatalf("Could not start server: %v", err)
	}
	defer s.Stop()

	db := s.DB()
	if err := db.Put("aa", "1"); err != nil {
		panic(err)
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	fmt.Printf("aa=%s\n", result.ValueBytes())

	// Output:
	// aa=1
}

func TestDebugName(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	file, _, _ := caller.Lookup(0)
	if err := db.Txn(func(txn *client.Txn) error {
		if !strings.HasPrefix(txn.DebugName(), file+":") {
			t.Fatalf("expected \"%s\" to have the prefix \"%s:\"", txn.DebugName(), file)
		}
		return nil
	}); err != nil {
		t.Errorf("txn failed: %s", err)
	}
}

func TestCommonMethods(t *testing.T) {
	defer leaktest.AfterTest(t)()
	batchType := reflect.TypeOf(&client.Batch{})
	dbType := reflect.TypeOf(&client.DB{})
	txnType := reflect.TypeOf(&client.Txn{})
	types := []reflect.Type{batchType, dbType, txnType}

	type key struct {
		typ    reflect.Type
		method string
	}
	omittedChecks := map[key]struct{}{
		// TODO(tschottdorf): removed GetProto from Batch, which necessitates
		// these two exceptions. Batch.GetProto would require wrapping each
		// request with the information that this particular Get must be
		// unmarshaled, which didn't seem worth doing as we're not using
		// Batch.GetProto at the moment.
		key{dbType, "GetProto"}:                {},
		key{txnType, "GetProto"}:               {},
		key{batchType, "CheckConsistency"}:     {},
		key{batchType, "AddRawRequest"}:        {},
		key{batchType, "PutInline"}:            {},
		key{batchType, "RawResponse"}:          {},
		key{batchType, "MustPErr"}:             {},
		key{dbType, "AdminMerge"}:              {},
		key{dbType, "AdminSplit"}:              {},
		key{dbType, "CheckConsistency"}:        {},
		key{dbType, "NewBatch"}:                {},
		key{dbType, "Run"}:                     {},
		key{dbType, "Txn"}:                     {},
		key{dbType, "GetSender"}:               {},
		key{dbType, "PutInline"}:               {},
		key{txnType, "Commit"}:                 {},
		key{txnType, "CommitInBatch"}:          {},
		key{txnType, "CommitOrCleanup"}:        {},
		key{txnType, "Rollback"}:               {},
		key{txnType, "CleanupOnError"}:         {},
		key{txnType, "DebugName"}:              {},
		key{txnType, "InternalSetPriority"}:    {},
		key{txnType, "IsFinalized"}:            {},
		key{txnType, "NewBatch"}:               {},
		key{txnType, "Exec"}:                   {},
		key{txnType, "Run"}:                    {},
		key{txnType, "SetDebugName"}:           {},
		key{txnType, "SetIsolation"}:           {},
		key{txnType, "SetUserPriority"}:        {},
		key{txnType, "SetSystemConfigTrigger"}: {},
		key{txnType, "SystemConfigTrigger"}:    {},
		key{txnType, "UpdateDeadlineMaybe"}:    {},
	}

	for b := range omittedChecks {
		if _, ok := b.typ.MethodByName(b.method); !ok {
			t.Fatalf("blacklist method (%s).%s does not exist", b.typ, b.method)
		}
	}

	for _, typ := range types {
		for j := 0; j < typ.NumMethod(); j++ {
			m := typ.Method(j)
			if len(m.PkgPath) > 0 {
				continue
			}
			if _, ok := omittedChecks[key{typ, m.Name}]; ok {
				continue
			}
			for _, otherTyp := range types {
				if typ == otherTyp {
					continue
				}
				if _, ok := otherTyp.MethodByName(m.Name); !ok {
					t.Errorf("(%s).%s does not exist, but (%s).%s does",
						otherTyp, m.Name, typ, m.Name)
				}
			}
		}
	}
}
<commit_msg>convert client examples to tests<commit_after>// Copyright 2015 The Cockroach 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.
//
// Author: Peter Mattis (peter@cockroachlabs.com)

package client_test

import (
	"bytes"
	"reflect"
	"strings"
	"testing"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/util/caller"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/log"
)

func setup() (server.TestServer, *client.DB) {
	s := server.StartTestServer(nil)
	return s, s.DB()
}

func checkIntResult(t *testing.T, expected, result int64) {
	if expected != result {
		t.Errorf("expected %d, got %d", expected, result)
	}
}

func checkResult(t *testing.T, expected, result []byte) {
	if !bytes.Equal(expected, result) {
		t.Errorf("expected \"%s\", got \"%s\"", expected, result)
	}
}

func checkResults(t *testing.T, expected map[string][]byte, results []client.Result) {
	count := 0
	for _, result := range results {
		checkRows(t, expected, result.Rows)
		count++
	}
	checkLen(t, len(expected), count)
}

func checkRows(t *testing.T, expected map[string][]byte, rows []client.KeyValue) {
	for i, row := range rows {
		if !bytes.Equal(expected[string(row.Key)], row.ValueBytes()) {
			t.Errorf("expected %d: %s=\"%s\", got %s=\"%s\"",
				i,
				row.Key,
				expected[string(row.Key)],
				row.Key,
				row.ValueBytes())
		}
	}
}

func checkLen(t *testing.T, expected, count int) {
	if expected != count {
		t.Errorf("expected length to be %d, got %d", expected, count)
	}
}

func TestDB_Get(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	checkResult(t, []byte(""), result.ValueBytes())
}

func TestDB_Put(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	if err := db.Put("aa", "1"); err != nil {
		panic(err)
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	checkResult(t, []byte("1"), result.ValueBytes())
}

func TestDB_CPut(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	if err := db.Put("aa", "1"); err != nil {
		panic(err)
	}
	if err := db.CPut("aa", "2", "1"); err != nil {
		panic(err)
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	checkResult(t, []byte("2"), result.ValueBytes())

	if err = db.CPut("aa", "3", "1"); err == nil {
		panic("expected error from conditional put")
	}
	result, err = db.Get("aa")
	if err != nil {
		panic(err)
	}
	checkResult(t, []byte("2"), result.ValueBytes())

	if err = db.CPut("bb", "4", "1"); err == nil {
		panic("expected error from conditional put")
	}
	result, err = db.Get("bb")
	if err != nil {
		panic(err)
	}
	checkResult(t, []byte(""), result.ValueBytes())

	if err = db.CPut("bb", "4", nil); err != nil {
		panic(err)
	}
	result, err = db.Get("bb")
	if err != nil {
		panic(err)
	}
	checkResult(t, []byte("4"), result.ValueBytes())
}

func TestDB_InitPut(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	if err := db.InitPut("aa", "1"); err != nil {
		panic(err)
	}
	if err := db.InitPut("aa", "1"); err != nil {
		panic(err)
	}
	if err := db.InitPut("aa", "2"); err == nil {
		panic("expected error from init put")
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	checkResult(t, []byte("1"), result.ValueBytes())
}

func TestDB_Inc(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	if _, err := db.Inc("aa", 100); err != nil {
		panic(err)
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	checkIntResult(t, 100, result.ValueInt())
}

func TestBatch(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	b := &client.Batch{}
	b.Get("aa")
	b.Put("bb", "2")
	if err := db.Run(b); err != nil {
		panic(err)
	}

	expected := map[string][]byte{
		"aa": []byte(""),
		"bb": []byte("2"),
	}
	checkResults(t, expected, b.Results)
}

func TestDB_Scan(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	b := &client.Batch{}
	b.Put("aa", "1")
	b.Put("ab", "2")
	b.Put("bb", "3")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	rows, err := db.Scan("a", "b", 100)
	if err != nil {
		panic(err)
	}
	expected := map[string][]byte{
		"aa": []byte("1"),
		"ab": []byte("2"),
	}

	checkRows(t, expected, rows)
	checkLen(t, len(expected), len(rows))
}

func TestDB_ReverseScan(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	b := &client.Batch{}
	b.Put("aa", "1")
	b.Put("ab", "2")
	b.Put("bb", "3")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	rows, err := db.ReverseScan("ab", "c", 100)
	if err != nil {
		panic(err)
	}
	expected := map[string][]byte{
		"bb": []byte("3"),
		"ab": []byte("2"),
	}

	checkRows(t, expected, rows)
	checkLen(t, len(expected), len(rows))
}

func TestDB_Del(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	b := &client.Batch{}
	b.Put("aa", "1")
	b.Put("ab", "2")
	b.Put("ac", "3")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	if err := db.Del("ab"); err != nil {
		panic(err)
	}
	rows, err := db.Scan("a", "b", 100)
	if err != nil {
		panic(err)
	}
	expected := map[string][]byte{
		"aa": []byte("1"),
		"ac": []byte("3"),
	}
	checkRows(t, expected, rows)
	checkLen(t, len(expected), len(rows))
}

func TestTxn_Commit(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	err := db.Txn(func(txn *client.Txn) error {
		b := txn.NewBatch()
		b.Put("aa", "1")
		b.Put("ab", "2")
		return txn.CommitInBatch(b)
	})
	if err != nil {
		panic(err)
	}

	b := &client.Batch{}
	b.Get("aa")
	b.Get("ab")
	if err := db.Run(b); err != nil {
		panic(err)
	}
	expected := map[string][]byte{
		"aa": []byte("1"),
		"ab": []byte("2"),
	}
	checkResults(t, expected, b.Results)
}

func TestDB_Put_insecure(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ctx := server.MakeTestContext()
	ctx.Insecure = true
	s := server.TestServer{
		Ctx: &ctx,
	}
	if err := s.Start(); err != nil {
		log.Fatalf("Could not start server: %v", err)
	}
	defer s.Stop()

	db := s.DB()
	if err := db.Put("aa", "1"); err != nil {
		panic(err)
	}
	result, err := db.Get("aa")
	if err != nil {
		panic(err)
	}
	checkResult(t, []byte("1"), result.ValueBytes())
}

func TestDebugName(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s, db := setup()
	defer s.Stop()

	file, _, _ := caller.Lookup(0)
	if err := db.Txn(func(txn *client.Txn) error {
		if !strings.HasPrefix(txn.DebugName(), file+":") {
			t.Fatalf("expected \"%s\" to have the prefix \"%s:\"", txn.DebugName(), file)
		}
		return nil
	}); err != nil {
		t.Errorf("txn failed: %s", err)
	}
}

func TestCommonMethods(t *testing.T) {
	defer leaktest.AfterTest(t)()
	batchType := reflect.TypeOf(&client.Batch{})
	dbType := reflect.TypeOf(&client.DB{})
	txnType := reflect.TypeOf(&client.Txn{})
	types := []reflect.Type{batchType, dbType, txnType}

	type key struct {
		typ    reflect.Type
		method string
	}
	omittedChecks := map[key]struct{}{
		// TODO(tschottdorf): removed GetProto from Batch, which necessitates
		// these two exceptions. Batch.GetProto would require wrapping each
		// request with the information that this particular Get must be
		// unmarshaled, which didn't seem worth doing as we're not using
		// Batch.GetProto at the moment.
		key{dbType, "GetProto"}:                {},
		key{txnType, "GetProto"}:               {},
		key{batchType, "CheckConsistency"}:     {},
		key{batchType, "AddRawRequest"}:        {},
		key{batchType, "PutInline"}:            {},
		key{batchType, "RawResponse"}:          {},
		key{batchType, "MustPErr"}:             {},
		key{dbType, "AdminMerge"}:              {},
		key{dbType, "AdminSplit"}:              {},
		key{dbType, "CheckConsistency"}:        {},
		key{dbType, "NewBatch"}:                {},
		key{dbType, "Run"}:                     {},
		key{dbType, "Txn"}:                     {},
		key{dbType, "GetSender"}:               {},
		key{dbType, "PutInline"}:               {},
		key{txnType, "Commit"}:                 {},
		key{txnType, "CommitInBatch"}:          {},
		key{txnType, "CommitOrCleanup"}:        {},
		key{txnType, "Rollback"}:               {},
		key{txnType, "CleanupOnError"}:         {},
		key{txnType, "DebugName"}:              {},
		key{txnType, "InternalSetPriority"}:    {},
		key{txnType, "IsFinalized"}:            {},
		key{txnType, "NewBatch"}:               {},
		key{txnType, "Exec"}:                   {},
		key{txnType, "Run"}:                    {},
		key{txnType, "SetDebugName"}:           {},
		key{txnType, "SetIsolation"}:           {},
		key{txnType, "SetUserPriority"}:        {},
		key{txnType, "SetSystemConfigTrigger"}: {},
		key{txnType, "SystemConfigTrigger"}:    {},
		key{txnType, "UpdateDeadlineMaybe"}:    {},
	}

	for b := range omittedChecks {
		if _, ok := b.typ.MethodByName(b.method); !ok {
			t.Fatalf("blacklist method (%s).%s does not exist", b.typ, b.method)
		}
	}

	for _, typ := range types {
		for j := 0; j < typ.NumMethod(); j++ {
			m := typ.Method(j)
			if len(m.PkgPath) > 0 {
				continue
			}
			if _, ok := omittedChecks[key{typ, m.Name}]; ok {
				continue
			}
			for _, otherTyp := range types {
				if typ == otherTyp {
					continue
				}
				if _, ok := otherTyp.MethodByName(m.Name); !ok {
					t.Errorf("(%s).%s does not exist, but (%s).%s does",
						otherTyp, m.Name, typ, m.Name)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>expr: Fix case where ParseExprDisjunct and ParseExprConjunct would panic an array index error.<commit_after><|endoftext|>


<commit_before>// Copyright 2018 Sean.ZH

package clients

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
)

// DnspodClient defines a client
type DnspodClient struct {
	Token string
}

// NewDnspodClient returns a new client
func NewDnspodClient(token string) *DnspodClient {
	return &DnspodClient{Token: token}
}

// ErrBadStatus is a error status for dnspod
var ErrBadStatus = errors.New("status is not 1")

// Status we just need to known it's status
type Status struct {
	Code string `json:"code"`
	// ...
}

// DNSPodRecordModify modify a record
func dNSPodRecordModify(token, domain, sub, rid, nip string) error {
	type RecordModifyStruct struct {
		Status `json:"status"`
	}
	v := url.Values{
		"domain":         {domain},
		"record_id":      {rid},
		"sub_domain":     {sub},
		"record_type":    {"A"},
		"record_line":    {"默认"},
		"record_line_id": {"0"},
		"value":          {nip},
		"ttl":            {"600"},
		"format":         {"json"},
		"login_token":    {token},
	}
	u := "https://dnsapi.cn/Record.Modify"
	var dpr RecordModifyStruct
	err := sendPost(u, &v, &dpr)
	if err != nil {
		return err
	}
	if dpr.Status.Code != "1" {
		fmt.Println("status code is", dpr)
		return ErrBadStatus
	}
	return nil
}

// DNSPodRecordList read all records
func dNSPodRecordList(token, domain, sub string) (string, error) {
	type RecordStruct struct {
		Id string `json:"id"`
	}
	type RecordListStruct struct {
		Status  `json:"status"`
		Records []RecordStruct `json:"records"`
	}
	v := url.Values{
		"domain":      {domain},
		"sub_domain":  {sub},
		"login_token": {token},
		"format":      {"json"},
	}
	u := "https://dnsapi.cn/Record.List"
	var dpr RecordListStruct
	err := sendPost(u, &v, &dpr)
	if err != nil {
		return "", err
	}
	if dpr.Status.Code != "1" {
		fmt.Println(dpr.Status.Code)
		return "", ErrBadStatus
	}
	return dpr.Records[0].Id, nil
}

// SendPost send post to api
func sendPost(u string, v *url.Values, ret interface{}) error {
	resp, err := http.PostForm(u, *v)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	bt, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(bt, &ret)
}

// ModifyRecord is a demo
func (d *DnspodClient) ModifyRecord(sub, domain, nip string) error {
	rid, err := dNSPodRecordList(d.Token, domain, sub)
	if err != nil {
		return err
	}
	return dNSPodRecordModify(d.Token, domain, sub, rid, nip)
}

// Call ModifyRecord(sub, domain, nip)
<commit_msg>change fmt to log<commit_after>// Copyright 2018 Sean.ZH

package clients

import (
	"encoding/json"
	"errors"
	"log"
	"io"
	"net/http"
	"net/url"
)

// DnspodClient defines a client
type DnspodClient struct {
	Token string
}

// NewDnspodClient returns a new client
func NewDnspodClient(token string) *DnspodClient {
	return &DnspodClient{Token: token}
}

// ErrBadStatus is a error status for dnspod
var ErrBadStatus = errors.New("status is not 1")

// Status we just need to known it's status
type Status struct {
	Code string `json:"code"`
	// ...
}

// DNSPodRecordModify modify a record
func dNSPodRecordModify(token, domain, sub, rid, nip string) error {
	type RecordModifyStruct struct {
		Status `json:"status"`
	}
	v := url.Values{
		"domain":         {domain},
		"record_id":      {rid},
		"sub_domain":     {sub},
		"record_type":    {"A"},
		"record_line":    {"默认"},
		"record_line_id": {"0"},
		"value":          {nip},
		"ttl":            {"600"},
		"format":         {"json"},
		"login_token":    {token},
	}
	u := "https://dnsapi.cn/Record.Modify"
	var dpr RecordModifyStruct
	err := sendPost(u, &v, &dpr)
	if err != nil {
		return err
	}
	if dpr.Status.Code != "1" {
		log.Println("status code is", dpr)
		return ErrBadStatus
	}
	return nil
}

// DNSPodRecordList read all records
func dNSPodRecordList(token, domain, sub string) (string, error) {
	type RecordStruct struct {
		Id string `json:"id"`
	}
	type RecordListStruct struct {
		Status  `json:"status"`
		Records []RecordStruct `json:"records"`
	}
	v := url.Values{
		"domain":      {domain},
		"sub_domain":  {sub},
		"login_token": {token},
		"format":      {"json"},
	}
	u := "https://dnsapi.cn/Record.List"
	var dpr RecordListStruct
	err := sendPost(u, &v, &dpr)
	if err != nil {
		return "", err
	}
	if dpr.Status.Code != "1" {
		log.Println(dpr.Status.Code)
		return "", ErrBadStatus
	}
	return dpr.Records[0].Id, nil
}

// SendPost send post to api
func sendPost(u string, v *url.Values, ret interface{}) error {
	resp, err := http.PostForm(u, *v)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	bt, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(bt, &ret)
}

// ModifyRecord is a demo
func (d *DnspodClient) ModifyRecord(sub, domain, nip string) error {
	rid, err := dNSPodRecordList(d.Token, domain, sub)
	if err != nil {
		return err
	}
	return dNSPodRecordModify(d.Token, domain, sub, rid, nip)
}

// Call ModifyRecord(sub, domain, nip)
<|endoftext|>


<commit_before>package node

import (
	// "reflect"
	"strconv"
	"strings"
)

// A Dealer routes and manages RPC calls to callees.
type Dealer interface {
	// Register a procedure on an endpoint
	Register(Sender, *Register)
	// Unregister a procedure on an endpoint
	Unregister(Sender, *Unregister)
	// Call a procedure on an endpoint
	Call(Sender, *Call)
	// Return the result of a procedure call
	Yield(Sender, *Yield)
	// Handle an ERROR message from an invocation
	Error(Sender, *Error)
	dump() string
	hasRegistration(string) bool
	lostSession(*Session)
}

type RemoteProcedure struct {
	Endpoint  Sender
	Procedure URI
	PassDetails bool
}

func NewRemoteProcedure(endpoint Sender, procedure URI, tags []string) RemoteProcedure {
	proc := RemoteProcedure{
		Endpoint: endpoint,
		Procedure: procedure,
		PassDetails: false,
	}

	for _, tag := range tags {
		switch {
		case tag == "details":
			proc.PassDetails = true
		}
	}

	return proc
}

type defaultDealer struct {
	// map registration IDs to procedures
	procedures map[ID]RemoteProcedure
	// map procedure URIs to registration IDs
	registrations map[URI]ID
	// keep track of call IDs so we can send the response to the caller
	calls map[ID]Sender
	// link the invocation ID to the call ID
	invocations map[ID]ID

	// Keep track of registrations by session, so that we can clean up when the
	// session closes.  For each session, we have a map[URI]bool, which we are
	// using as a set of registrations (store true for register, delete for
	// unregister).
	sessionRegistrations map[Sender]map[URI]bool
}

func NewDefaultDealer() Dealer {
	return &defaultDealer{
		procedures:    make(map[ID]RemoteProcedure),
		registrations: make(map[URI]ID),
		calls:         make(map[ID]Sender),
		invocations:   make(map[ID]ID),
		sessionRegistrations: make(map[Sender]map[URI]bool),
	}
}

func (d *defaultDealer) Register(callee Sender, msg *Register) {
	// Endpoint may contain a # sign to pass comma-separated tags.
	// Example: pd.agent/function#details
	parts := strings.SplitN(string(msg.Procedure), "#", 2)
	endpoint := URI(parts[0])

	var tags []string
	if len(parts) > 1 {
		tags = strings.Split(parts[1], ",")
	}

	if id, ok := d.registrations[endpoint]; ok {
		//log.Println("error: procedure already exists:", msg.Procedure, id)
		out.Error("error: procedure already exists:", endpoint, id)
		callee.Send(&Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Details: make(map[string]interface{}),
			Error:   ErrProcedureAlreadyExists,
		})
		return
	}
	reg := NewID()
	d.procedures[reg] = NewRemoteProcedure(callee, endpoint, tags)
	d.registrations[endpoint] = reg

	if d.sessionRegistrations[callee] == nil {
		d.sessionRegistrations[callee] = make(map[URI]bool)
	}
	d.sessionRegistrations[callee][endpoint] = true

	//log.Printf("registered procedure %v [%v]", reg, msg.Procedure)
	callee.Send(&Registered{
		Request:      msg.Request,
		Registration: reg,
	})
}

func (d *defaultDealer) Unregister(callee Sender, msg *Unregister) {
	if procedure, ok := d.procedures[msg.Registration]; !ok {
		// the registration doesn't exist
		//log.Println("error: no such registration:", msg.Registration)
		callee.Send(&Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Details: make(map[string]interface{}),
			Error:   ErrNoSuchRegistration,
		})
	} else {
		delete(d.sessionRegistrations[callee], procedure.Procedure)
		delete(d.registrations, procedure.Procedure)
		delete(d.procedures, msg.Registration)
		//log.Printf("unregistered procedure %v [%v]", procedure.Procedure, msg.Registration)
		callee.Send(&Unregistered{
			Request: msg.Request,
		})
	}
}

func (d *defaultDealer) Call(caller Sender, msg *Call) {
	if reg, ok := d.registrations[msg.Procedure]; !ok {
		caller.Send(&Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Details: make(map[string]interface{}),
			Error:   ErrNoSuchProcedure,
		})
	} else {
		if rproc, ok := d.procedures[reg]; !ok {
			// found a registration id, but doesn't match any remote procedure
			caller.Send(&Error{
				Type:    msg.MessageType(),
				Request: msg.Request,
				Details: make(map[string]interface{}),
				// TODO: what should this error be?
				Error: URI("wamp.error.internal_error"),
			})
		} else {
			// everything checks out, make the invocation request
			// TODO: make the Request ID specific to the caller

			args := msg.Arguments
			kwargs := msg.ArgumentsKw

			// Remote procedures with the PassDetails flag set will receive a
			// special first argument set by the node.
			if rproc.PassDetails {
				details := make(map[string]interface{})

				// Make sure the argument list exists first.
				if args == nil {
					args = make([]interface{}, 0)
				}

				// Does the caller want to be disclosed?
				// We default to true unless he explicitly says otherwise.
				disclose_caller, ok := msg.Options["disclose_me"].(bool)
				if !ok {
					disclose_caller = true
				}

				if disclose_caller {
					sess := caller.(*Session)
					if sess != nil {
						details["caller"] = sess.pdid
					}
				}

				// Insert as the first positional argument.
				args = append(args, nil)
				copy(args[1:], args[:])
				args[0] = details
			}

			d.calls[msg.Request] = caller
			invocationID := NewID()
			d.invocations[invocationID] = msg.Request
			rproc.Endpoint.Send(&Invocation{
				Request:      invocationID,
				Registration: reg,
				Details:      map[string]interface{}{},
				Arguments:    args,
				ArgumentsKw:  kwargs,
			})
		}
	}
}

func (d *defaultDealer) Yield(callee Sender, msg *Yield) {
	if callID, ok := d.invocations[msg.Request]; !ok {
		// WAMP spec doesn't allow sending an error in response to a YIELD message
		//log.Println("received YIELD message with invalid invocation request ID:", msg.Request)
	} else {
		delete(d.invocations, msg.Request)
		if caller, ok := d.calls[callID]; !ok {
			// found the invocation id, but doesn't match any call id
			// WAMP spec doesn't allow sending an error in response to a YIELD message
			//log.Printf("received YIELD message, but unable to match it (%v) to a CALL ID", msg.Request)
		} else {
			delete(d.calls, callID)
			// return the result to the caller
			caller.Send(&Result{
				Request:     callID,
				Details:     map[string]interface{}{},
				Arguments:   msg.Arguments,
				ArgumentsKw: msg.ArgumentsKw,
			})
			//log.Printf("returned YIELD %v to caller as RESULT %v", msg.Request, callID)
		}
	}
}

func (d *defaultDealer) Error(peer Sender, msg *Error) {
	if callID, ok := d.invocations[msg.Request]; !ok {
		//log.Println("received ERROR (INVOCATION) message with invalid invocation request ID:", msg.Request)
	} else {
		delete(d.invocations, msg.Request)
		if caller, ok := d.calls[callID]; !ok {
			//log.Printf("received ERROR (INVOCATION) message, but unable to match it (%v) to a CALL ID", msg.Request)
		} else {
			delete(d.calls, callID)
			// return an error to the caller
			caller.Send(&Error{
				Type:        CALL,
				Request:     callID,
				Details:     make(map[string]interface{}),
				Arguments:   msg.Arguments,
				ArgumentsKw: msg.ArgumentsKw,
				Error:       msg.Error,
			})
			//log.Printf("returned ERROR %v to caller as ERROR %v", msg.Request, callID)
		}
	}
}

// Remove all the registrations for a session that has disconected
func (d *defaultDealer) lostSession(sess *Session) {
	for uri, _ := range(d.sessionRegistrations[sess]) {
		out.Debug("Unregister: %s", string(uri))
		delete(d.procedures, d.registrations[uri])
		delete(d.registrations, uri)
	}

	delete(d.sessionRegistrations, sess)
}

func (d *defaultDealer) dump() string {
	ret := "  functions:"

	for k, v := range d.procedures {
		ret += "\n\t" + strconv.FormatUint(uint64(k), 16) + ": " + string(v.Procedure)
	}

	ret += "\n  registrations:"

	for k, v := range d.registrations {
		ret += "\n\t" + string(k) + ": " + strconv.FormatUint(uint64(v), 16)
	}

	ret += "\n  calls:"

	for k, _ := range d.calls {
		ret += "\n\t" + strconv.FormatUint(uint64(k), 16) + ": (sender)"
	}

	ret += "\n  invocations:"

	for k, v := range d.invocations {
		ret += "\n\t" + strconv.FormatUint(uint64(k), 16) + ": " + strconv.FormatUint(uint64(v), 16)
	}

	return ret
}

// Testing. Not sure if this works 100 or not
func (d *defaultDealer) hasRegistration(s string) bool {
	_, exists := d.registrations[URI(s)]
	return exists
}
<commit_msg>Tracking calls by node-generated InvocationId rather than caller-supplied RequestId.<commit_after>package node

import (
	// "reflect"
	"strconv"
	"strings"
)

// A Dealer routes and manages RPC calls to callees.
type Dealer interface {
	// Register a procedure on an endpoint
	Register(Sender, *Register)
	// Unregister a procedure on an endpoint
	Unregister(Sender, *Unregister)
	// Call a procedure on an endpoint
	Call(Sender, *Call)
	// Return the result of a procedure call
	Yield(Sender, *Yield)
	// Handle an ERROR message from an invocation
	Error(Sender, *Error)
	dump() string
	hasRegistration(string) bool
	lostSession(*Session)
}

type RemoteProcedure struct {
	Endpoint  Sender
	Procedure URI
	PassDetails bool
}

func NewRemoteProcedure(endpoint Sender, procedure URI, tags []string) RemoteProcedure {
	proc := RemoteProcedure{
		Endpoint: endpoint,
		Procedure: procedure,
		PassDetails: false,
	}

	for _, tag := range tags {
		switch {
		case tag == "details":
			proc.PassDetails = true
		}
	}

	return proc
}

type defaultDealer struct {
	// map registration IDs to procedures
	procedures map[ID]RemoteProcedure
	// map procedure URIs to registration IDs
	registrations map[URI]ID

	// Map InvocationID to RequestID so we can send the RequestID with the
	// result (lets caller know what request the result is for).
	requests map[ID]ID

	// Map InvocationID to Sender so we know where to send the response.
	callers map[ID]Sender

	// Keep track of registrations by session, so that we can clean up when the
	// session closes.  For each session, we have a map[URI]bool, which we are
	// using as a set of registrations (store true for register, delete for
	// unregister).
	sessionRegistrations map[Sender]map[URI]bool
}

func NewDefaultDealer() Dealer {
	return &defaultDealer{
		procedures:    make(map[ID]RemoteProcedure),
		registrations: make(map[URI]ID),
		requests:      make(map[ID]ID),
		callers:       make(map[ID]Sender),
		sessionRegistrations: make(map[Sender]map[URI]bool),
	}
}

func (d *defaultDealer) Register(callee Sender, msg *Register) {
	// Endpoint may contain a # sign to pass comma-separated tags.
	// Example: pd.agent/function#details
	parts := strings.SplitN(string(msg.Procedure), "#", 2)
	endpoint := URI(parts[0])

	var tags []string
	if len(parts) > 1 {
		tags = strings.Split(parts[1], ",")
	}

	if id, ok := d.registrations[endpoint]; ok {
		//log.Println("error: procedure already exists:", msg.Procedure, id)
		out.Error("error: procedure already exists:", endpoint, id)
		callee.Send(&Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Details: make(map[string]interface{}),
			Error:   ErrProcedureAlreadyExists,
		})
		return
	}

	reg := NewID()
	d.procedures[reg] = NewRemoteProcedure(callee, endpoint, tags)
	d.registrations[endpoint] = reg

	if d.sessionRegistrations[callee] == nil {
		d.sessionRegistrations[callee] = make(map[URI]bool)
	}
	d.sessionRegistrations[callee][endpoint] = true

	//log.Printf("registered procedure %v [%v]", reg, msg.Procedure)
	callee.Send(&Registered{
		Request:      msg.Request,
		Registration: reg,
	})
}

func (d *defaultDealer) Unregister(callee Sender, msg *Unregister) {
	if procedure, ok := d.procedures[msg.Registration]; !ok {
		// the registration doesn't exist
		//log.Println("error: no such registration:", msg.Registration)
		callee.Send(&Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Details: make(map[string]interface{}),
			Error:   ErrNoSuchRegistration,
		})
	} else {
		delete(d.sessionRegistrations[callee], procedure.Procedure)
		delete(d.registrations, procedure.Procedure)
		delete(d.procedures, msg.Registration)
		//log.Printf("unregistered procedure %v [%v]", procedure.Procedure, msg.Registration)
		callee.Send(&Unregistered{
			Request: msg.Request,
		})
	}
}

func (d *defaultDealer) Call(caller Sender, msg *Call) {
	if reg, ok := d.registrations[msg.Procedure]; !ok {
		caller.Send(&Error{
			Type:    msg.MessageType(),
			Request: msg.Request,
			Details: make(map[string]interface{}),
			Error:   ErrNoSuchProcedure,
		})
	} else {
		if rproc, ok := d.procedures[reg]; !ok {
			// found a registration id, but doesn't match any remote procedure
			caller.Send(&Error{
				Type:    msg.MessageType(),
				Request: msg.Request,
				Details: make(map[string]interface{}),
				// TODO: what should this error be?
				Error: URI("wamp.error.internal_error"),
			})
		} else {
			// everything checks out, make the invocation request
			args := msg.Arguments
			kwargs := msg.ArgumentsKw

			// Remote procedures with the PassDetails flag set will receive a
			// special first argument set by the node.
			if rproc.PassDetails {
				details := make(map[string]interface{})

				// Make sure the argument list exists first.
				if args == nil {
					args = make([]interface{}, 0)
				}

				// Does the caller want to be disclosed?
				// We default to true unless he explicitly says otherwise.
				disclose_caller, ok := msg.Options["disclose_me"].(bool)
				if !ok {
					disclose_caller = true
				}

				if disclose_caller {
					sess := caller.(*Session)
					if sess != nil {
						details["caller"] = sess.pdid
					}
				}

				// Insert as the first positional argument.
				args = append(args, nil)
				copy(args[1:], args[:])
				args[0] = details
			}

			invocationID := NewID()
			d.requests[invocationID] = msg.Request
			d.callers[invocationID] = caller

			rproc.Endpoint.Send(&Invocation{
				Request:      invocationID,
				Registration: reg,
				Details:      map[string]interface{}{},
				Arguments:    args,
				ArgumentsKw:  kwargs,
			})
		}
	}
}

func (d *defaultDealer) Yield(callee Sender, msg *Yield) {
	caller, ok := d.callers[msg.Request]
	if !ok {
		// WAMP spec doesn't allow sending an error in response to a YIELD message
		//log.Println("received YIELD message with invalid invocation request ID:", msg.Request)
		return
	}

	delete(d.callers, msg.Request)

	requestId, ok := d.requests[msg.Request]
	if !ok {
		return
	}

	delete(d.requests, msg.Request)

	// return the result to the caller
	caller.Send(&Result{
		Request:     requestId,
		Details:     map[string]interface{}{},
		Arguments:   msg.Arguments,
		ArgumentsKw: msg.ArgumentsKw,
	})
}

func (d *defaultDealer) Error(peer Sender, msg *Error) {
	caller, ok := d.callers[msg.Request]
	if !ok {
		//log.Println("received ERROR (INVOCATION) message with invalid invocation request ID:", msg.Request)
		return
	}

	delete(d.callers, msg.Request)

	requestId, ok := d.requests[msg.Request]
	if !ok {
		//log.Printf("received ERROR (INVOCATION) message, but unable to match it (%v) to a CALL ID", msg.Request)
		return
	}

	delete(d.requests, msg.Request)

	// return an error to the caller
	caller.Send(&Error{
		Type:        CALL,
		Request:     requestId,
		Details:     make(map[string]interface{}),
		Arguments:   msg.Arguments,
		ArgumentsKw: msg.ArgumentsKw,
		Error:       msg.Error,
	})
}

// Remove all the registrations for a session that has disconected
func (d *defaultDealer) lostSession(sess *Session) {
	// TODO: Do something about outstanding requests

	for uri, _ := range(d.sessionRegistrations[sess]) {
		out.Debug("Unregister: %s", string(uri))
		delete(d.procedures, d.registrations[uri])
		delete(d.registrations, uri)
	}

	delete(d.sessionRegistrations, sess)
}

func (d *defaultDealer) dump() string {
	ret := "  functions:"

	for k, v := range d.procedures {
		ret += "\n\t" + strconv.FormatUint(uint64(k), 16) + ": " + string(v.Procedure)
	}

	ret += "\n  registrations:"

	for k, v := range d.registrations {
		ret += "\n\t" + string(k) + ": " + strconv.FormatUint(uint64(v), 16)
	}

	ret += "\n  callers:"

	for k, _ := range d.callers {
		ret += "\n\t" + strconv.FormatUint(uint64(k), 16) + ": (sender)"
	}

	ret += "\n  requests:"

	for k, v := range d.requests {
		ret += "\n\t" + strconv.FormatUint(uint64(k), 16) + ": " + strconv.FormatUint(uint64(v), 16)
	}

	return ret
}

// Testing. Not sure if this works 100 or not
func (d *defaultDealer) hasRegistration(s string) bool {
	_, exists := d.registrations[URI(s)]
	return exists
}
<|endoftext|>