<commit_before>package aws

import (
	"os"
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/sns"
)

/**
 Before running this test, a few ENV variables must be set:
 GCM_API_KEY - Google Cloud Messaging API Key
 APNS_SANDBOX_CREDENTIAL - Apple Push Notification Sandbox Private Key
 APNS_SANDBOX_PRINCIPAL - Apple Push Notification Sandbox Certificate
**/

func TestAccAWSSNSApplication_gcm_create_update(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck: func() {
			testAccPreCheck(t)
			if os.Getenv("GCM_API_KEY") == "" {
				t.Fatal("GCM_API_KEY must be set.")
			}
		},
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSApplicationDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSSNSApplicationGCMConfig,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "name", "aws_sns_application_test"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "platform", "GCM"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "success_feedback_sample_rate", "100"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "event_endpoint_created_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-created-topic"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "event_endpoint_updated_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-updated-topic"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "event_delivery_failure_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-failure-topic"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "event_endpoint_deleted_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-deleted-topic"),
				),
			},
			resource.TestStep{
				Config: testAccAWSSNSApplicationGCMConfigUpdate,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "name", "aws_sns_application_test"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "platform", "GCM"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "success_feedback_sample_rate", "99"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "event_endpoint_created_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-created-topic-update"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "event_endpoint_updated_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-updated-topic-update"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "event_delivery_failure_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-failure-topic-update"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.gcm_test", "event_endpoint_deleted_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-deleted-topic-update"),
				),
			},
		},
	})
}

func TestAccAWSSNSApplication_apns_sandbox_create_update(t *testing.T) {

	resource.Test(t, resource.TestCase{
		PreCheck: func() {
			testAccPreCheck(t)
			if os.Getenv("APNS_SANDBOX_CREDENTIAL") == "" {
				t.Fatal("APNS_SANDBOX_CREDENTIAL must be set.")
			}
			if os.Getenv("APNS_SANDBOX_PRINCIPAL") == "" {
				t.Fatal("APNS_SANDBOX_PRINCIPAL must be set.")
			}
		},
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSApplicationDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSSNSApplicationAPNSSandBoxConfig,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "name", "aws_sns_application_test"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "platform", "APNS_SANDBOX"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "success_feedback_sample_rate", "100"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "event_endpoint_created_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-created-topic"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "event_endpoint_updated_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-updated-topic"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "event_delivery_failure_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-failure-topic"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "event_endpoint_deleted_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-deleted-topic"),
				),
			},
			resource.TestStep{
				Config: testAccAWSSNSApplicationAPNSSandBoxConfigUpdate,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "name", "aws_sns_application_test"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "platform", "APNS_SANDBOX"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "success_feedback_sample_rate", "99"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "event_endpoint_created_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-created-topic-update"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "event_endpoint_updated_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-updated-topic-update"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "event_delivery_failure_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-failure-topic-update"),
					resource.TestCheckResourceAttr(
						"aws_sns_application.apns_test", "event_endpoint_deleted_topic_arn", "arn:aws:sns:us-east-1:638386993804:endpoint-deleted-topic-update"),
				),
			},
		},
	})
}

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

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_sns_application" {
			continue
		}
		_, err := conn.DeletePlatformApplication(&sns.DeletePlatformApplicationInput{
			PlatformApplicationArn: aws.String(rs.Primary.ID),
		})
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "NoSNSApplication" {
				return nil
			}
			return err
		}
	}
	return nil
}

var testAccAWSSNSApplicationGCMConfig = `
resource "aws_sns_application" "gcm_test" {
	event_delivery_failure_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-failure-topic"
	event_endpoint_created_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-created-topic"
	event_endpoint_deleted_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-deleted-topic"
	event_endpoint_updated_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-updated-topic"
	name = "aws_sns_application_test"
	platform = "GCM"
	platform_credential = "` + os.Getenv("GCM_API_KEY") + `"
	success_feedback_sample_rate = 100
}
`

var testAccAWSSNSApplicationGCMConfigUpdate = `
	event_delivery_failure_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-failure-topic-update"
	event_endpoint_created_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-created-topic-update"
	event_endpoint_deleted_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-deleted-topic-update"
	event_endpoint_updated_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-updated-topic-update"
	name = "aws_sns_application_test"
	platform = "GCM"
	platform_credential = "` + os.Getenv("GCM_API_KEY") + `"
	resource "aws_sns_application" "gcm_test" {
	success_feedback_sample_rate = 99
}
`

var testAccAWSSNSApplicationAPNSSandBoxConfig = `
resource "aws_sns_application" "apns_test" {
	event_delivery_failure_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-failure-topic"
	event_endpoint_created_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-created-topic"
	event_endpoint_deleted_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-deleted-topic"
	event_endpoint_updated_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-updated-topic"
	name = "aws_sns_application_test"
	platform = "APNS_SANDBOX"
	platform_credential = "` + os.Getenv("APNS_SANDBOX_CREDENTIAL") + `"
	platform_principal = "` + os.Getenv("APNS_SANDBOX_PRINCIPAL") + `"
	success_feedback_sample_rate = 100
}
`

var testAccAWSSNSApplicationAPNSSandBoxConfigUpdate = `
resource "aws_sns_application" "apns_test" {
	event_delivery_failure_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-failure-topic-update"
	event_endpoint_created_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-created-topic-update"
	event_endpoint_deleted_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-deleted-topic-update"
	event_endpoint_updated_topic_arn = "arn:aws:sns:us-east-1:638386993804:endpoint-updated-topic-update"
	name = "aws_sns_application_test"
	platform = "APNS_SANDBOX"
	platform_credential = "` + os.Getenv("APNS_SANDBOX_CREDENTIAL") + `"
	platform_principal = "` + os.Getenv("APNS_SANDBOX_PRINCIPAL") + `"
	success_feedback_sample_rate = 99
}
`
<commit_msg>test/resource/aws_sns_application: Refactor testing<commit_after>package aws

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"testing"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	homedir "github.com/mitchellh/go-homedir"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sns"
)

/**
 Before running this test, at least one of these ENV variables combinations must be set:

 GCM_API_KEY - Google Cloud Messaging API Key

 APNS_SANDBOX_CREDENTIAL_PATH - Apple Push Notification Sandbox Private Key file location
 APNS_SANDBOX_PRINCIPAL_PATH - Apple Push Notification Sandbox Certificate file location
**/

type testAccAwsSnsApplicationPlatform struct {
	Name           string
	Credential     string
	CredentialHash string
	Principal      string
	PrincipalHash  string
}

func testAccAwsSnsApplicationPlatformFromEnv() ([]*testAccAwsSnsApplicationPlatform, error) {
	platforms := make([]*testAccAwsSnsApplicationPlatform, 0, 2)

	if os.Getenv("APNS_SANDBOX_CREDENTIAL_PATH") != "" {
		if os.Getenv("APNS_SANDBOX_PRINCIPAL_PATH") == "" {
			return platforms, errors.New("APNS_SANDBOX_CREDENTIAL_PATH set but missing APNS_SANDBOX_PRINCIPAL_PATH")
		}
		credentialHash, err := testAccHashSumPath(os.Getenv("APNS_SANDBOX_CREDENTIAL_PATH"))
		if err != nil {
			return platforms, err
		}
		principalHash, err := testAccHashSumPath(os.Getenv("APNS_SANDBOX_PRINCIPAL_PATH"))
		if err != nil {
			return platforms, err
		}

		platform := &testAccAwsSnsApplicationPlatform{
			Name:           "APNS_SANDBOX",
			Credential:     fmt.Sprintf("${file(pathexpand(%q))}", os.Getenv("APNS_SANDBOX_CREDENTIAL_PATH")),
			CredentialHash: credentialHash,
			Principal:      fmt.Sprintf("${file(pathexpand(%q))}", os.Getenv("APNS_SANDBOX_PRINCIPAL_PATH")),
			PrincipalHash:  principalHash,
		}
		platforms = append(platforms, platform)
	}

	if os.Getenv("GCM_API_KEY") != "" {
		platform := &testAccAwsSnsApplicationPlatform{
			Name:           "GCM",
			Credential:     os.Getenv("GCM_API_KEY"),
			CredentialHash: hashSum(os.Getenv("GCM_API_KEY")),
		}
		platforms = append(platforms, platform)
	}

	if len(platforms) == 0 {
		return platforms, errors.New("no SNS Platform Application environment variables found")
	}
	return platforms, nil
}

func testAccHashSumPath(path string) (string, error) {
	path, err := homedir.Expand(path)
	if err != nil {
		return "", err
	}
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return "", err
	}
	return hashSum(string(data)), nil
}

func TestAccAwsSnsApplication_basic(t *testing.T) {
	platforms, err := testAccAwsSnsApplicationPlatformFromEnv()
	if err != nil {
		t.Skip(err)
	}

	resourceName := "aws_sns_application.test"

	for _, platform := range platforms {
		name := fmt.Sprintf("tf-acc-%d", acctest.RandInt())
		platformPrincipalCheck := resource.TestCheckNoResourceAttr(resourceName, "platform_principal")
		if platform.Principal != "" {
			platformPrincipalCheck = resource.TestCheckResourceAttr(resourceName, "platform_principal", platform.PrincipalHash)
		}

		t.Run(platform.Name, func(*testing.T) {
			resource.Test(t, resource.TestCase{
				PreCheck:     func() { testAccPreCheck(t) },
				Providers:    testAccProviders,
				CheckDestroy: testAccCheckAWSSNSApplicationDestroy,
				Steps: []resource.TestStep{
					{
						Config: testAccAwsSnsApplicationConfig_basic(name, &testAccAwsSnsApplicationPlatform{
							Name:       "APNS",
							Credential: "NOTEMPTY",
							Principal:  "",
						}),
						ExpectError: regexp.MustCompile(`platform_principal is required when platform =`),
					},
					{
						Config: testAccAwsSnsApplicationConfig_basic(name, &testAccAwsSnsApplicationPlatform{
							Name:       "APNS_SANDBOX",
							Credential: "NOTEMPTY",
							Principal:  "",
						}),
						ExpectError: regexp.MustCompile(`platform_principal is required when platform =`),
					},
					{
						Config: testAccAwsSnsApplicationConfig_basic(name, platform),
						Check: resource.ComposeTestCheckFunc(
							testAccCheckAwsSnsApplicationExists(resourceName),
							resource.TestMatchResourceAttr(resourceName, "arn", regexp.MustCompile(fmt.Sprintf("^arn:[^:]+:sns:[^:]+:[^:]+:app/%s/%s$", platform.Name, name))),
							resource.TestCheckResourceAttr(resourceName, "name", name),
							resource.TestCheckResourceAttr(resourceName, "platform", platform.Name),
							resource.TestCheckResourceAttr(resourceName, "platform_credential", platform.CredentialHash),
							platformPrincipalCheck,
						),
					},
					{
						ResourceName:            resourceName,
						ImportState:             true,
						ImportStateVerify:       true,
						ImportStateVerifyIgnore: []string{"platform_credential", "platform_principal"},
					},
				},
			})
		})
	}
}

func TestAccAwsSnsApplication_basicAttributes(t *testing.T) {
	platforms, err := testAccAwsSnsApplicationPlatformFromEnv()
	if err != nil {
		t.Skip(err)
	}

	resourceName := "aws_sns_application.test"

	var testCases = []struct {
		AttributeKey         string
		AttributeValue       string
		AttributeValueUpdate string
	}{
		{
			AttributeKey:         "success_feedback_sample_rate",
			AttributeValue:       "100",
			AttributeValueUpdate: "99",
		},
	}

	for _, platform := range platforms {
		t.Run(platform.Name, func(*testing.T) {
			for _, tc := range testCases {
				t.Run(fmt.Sprintf("%s/%s", platform.Name, tc.AttributeKey), func(*testing.T) {
					name := fmt.Sprintf("tf-acc-%d", acctest.RandInt())

					resource.Test(t, resource.TestCase{
						PreCheck:     func() { testAccPreCheck(t) },
						Providers:    testAccProviders,
						CheckDestroy: testAccCheckAWSSNSApplicationDestroy,
						Steps: []resource.TestStep{
							{
								Config: testAccAwsSnsApplicationConfig_basicAttribute(name, platform, tc.AttributeKey, tc.AttributeValue),
								Check: resource.ComposeTestCheckFunc(
									testAccCheckAwsSnsApplicationExists(resourceName),
									resource.TestCheckResourceAttr(resourceName, tc.AttributeKey, tc.AttributeValue),
								),
							},
							{
								Config: testAccAwsSnsApplicationConfig_basicAttribute(name, platform, tc.AttributeKey, tc.AttributeValueUpdate),
								Check: resource.ComposeTestCheckFunc(
									testAccCheckAwsSnsApplicationExists(resourceName),
									resource.TestCheckResourceAttr(resourceName, tc.AttributeKey, tc.AttributeValueUpdate),
								),
							},
							{
								ResourceName:            resourceName,
								ImportState:             true,
								ImportStateVerify:       true,
								ImportStateVerifyIgnore: []string{"platform_credential", "platform_principal"},
							},
						},
					})
				})
			}
		})
	}
}

func TestAccAwsSnsApplication_iamRoleAttributes(t *testing.T) {
	platforms, err := testAccAwsSnsApplicationPlatformFromEnv()
	if err != nil {
		t.Skip(err)
	}

	resourceName := "aws_sns_application.test"

	var testCases = []string{
		"failure_feedback_role_arn",
		"success_feedback_role_arn",
	}

	for _, platform := range platforms {
		t.Run(platform.Name, func(*testing.T) {
			for _, tc := range testCases {
				t.Run(fmt.Sprintf("%s/%s", platform.Name, tc), func(*testing.T) {
					iamRoleName1 := fmt.Sprintf("tf-acc-%d", acctest.RandInt())
					iamRoleName2 := fmt.Sprintf("tf-acc-%d", acctest.RandInt())
					name := fmt.Sprintf("tf-acc-%d", acctest.RandInt())

					resource.Test(t, resource.TestCase{
						PreCheck:     func() { testAccPreCheck(t) },
						Providers:    testAccProviders,
						CheckDestroy: testAccCheckAWSSNSApplicationDestroy,
						Steps: []resource.TestStep{
							{
								Config: testAccAwsSnsApplicationConfig_iamRoleAttribute(name, platform, tc, iamRoleName1),
								Check: resource.ComposeTestCheckFunc(
									testAccCheckAwsSnsApplicationExists(resourceName),
									resource.TestMatchResourceAttr(resourceName, tc, regexp.MustCompile(fmt.Sprintf("^arn:[^:]+:iam::[^:]+:role/%s$", iamRoleName1))),
								),
							},
							{
								Config: testAccAwsSnsApplicationConfig_iamRoleAttribute(name, platform, tc, iamRoleName2),
								Check: resource.ComposeTestCheckFunc(
									testAccCheckAwsSnsApplicationExists(resourceName),
									resource.TestMatchResourceAttr(resourceName, tc, regexp.MustCompile(fmt.Sprintf("^arn:[^:]+:iam::[^:]+:role/%s$", iamRoleName2))),
								),
							},
							{
								ResourceName:            resourceName,
								ImportState:             true,
								ImportStateVerify:       true,
								ImportStateVerifyIgnore: []string{"platform_credential", "platform_principal"},
							},
						},
					})
				})
			}
		})
	}
}

func TestAccAwsSnsApplication_snsTopicAttributes(t *testing.T) {
	platforms, err := testAccAwsSnsApplicationPlatformFromEnv()
	if err != nil {
		t.Skip(err)
	}

	resourceName := "aws_sns_application.test"

	var testCases = []string{
		"event_delivery_failure_topic_arn",
		"event_endpoint_created_topic_arn",
		"event_endpoint_deleted_topic_arn",
		"event_endpoint_updated_topic_arn",
	}

	for _, platform := range platforms {
		t.Run(platform.Name, func(*testing.T) {
			for _, tc := range testCases {
				t.Run(fmt.Sprintf("%s/%s", platform.Name, tc), func(*testing.T) {
					snsTopicName1 := fmt.Sprintf("tf-acc-%d", acctest.RandInt())
					snsTopicName2 := fmt.Sprintf("tf-acc-%d", acctest.RandInt())
					name := fmt.Sprintf("tf-acc-%d", acctest.RandInt())

					resource.Test(t, resource.TestCase{
						PreCheck:     func() { testAccPreCheck(t) },
						Providers:    testAccProviders,
						CheckDestroy: testAccCheckAWSSNSApplicationDestroy,
						Steps: []resource.TestStep{
							{
								Config: testAccAwsSnsApplicationConfig_snsTopicAttribute(name, platform, tc, snsTopicName1),
								Check: resource.ComposeTestCheckFunc(
									testAccCheckAwsSnsApplicationExists(resourceName),
									resource.TestMatchResourceAttr(resourceName, tc, regexp.MustCompile(fmt.Sprintf("^arn:[^:]+:sns:[^:]+:[^:]+:%s$", snsTopicName1))),
								),
							},
							{
								Config: testAccAwsSnsApplicationConfig_snsTopicAttribute(name, platform, tc, snsTopicName2),
								Check: resource.ComposeTestCheckFunc(
									testAccCheckAwsSnsApplicationExists(resourceName),
									resource.TestMatchResourceAttr(resourceName, tc, regexp.MustCompile(fmt.Sprintf("^arn:[^:]+:sns:[^:]+:[^:]+:%s$", snsTopicName2))),
								),
							},
							{
								ResourceName:            resourceName,
								ImportState:             true,
								ImportStateVerify:       true,
								ImportStateVerifyIgnore: []string{"platform_credential", "platform_principal"},
							},
						},
					})
				})
			}
		})
	}
}

func testAccCheckAwsSnsApplicationExists(name 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("missing ID: %s", name)
		}

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

		input := &sns.GetPlatformApplicationAttributesInput{
			PlatformApplicationArn: aws.String(rs.Primary.ID),
		}

		log.Printf("[DEBUG] Reading SNS Platform Application attributes: %s", input)
		_, err := conn.GetPlatformApplicationAttributes(input)
		if err != nil {
			return err
		}

		return nil
	}
}

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

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

		input := &sns.GetPlatformApplicationAttributesInput{
			PlatformApplicationArn: aws.String(rs.Primary.ID),
		}

		log.Printf("[DEBUG] Reading SNS Platform Application attributes: %s", input)
		_, err := conn.GetPlatformApplicationAttributes(input)
		if err != nil {
			if isAWSErr(err, sns.ErrCodeNotFoundException, "") {
				return nil
			}
			return err
		}
	}
	return nil
}

func testAccAwsSnsApplicationConfig_basic(name string, platform *testAccAwsSnsApplicationPlatform) string {
	if platform.Principal == "" {
		return fmt.Sprintf(`
resource "aws_sns_application" "test" {
  name                = "%s"
  platform            = "%s"
  platform_credential = "%s"
}
`, name, platform.Name, platform.Credential)
	}
	return fmt.Sprintf(`
resource "aws_sns_application" "test" {
  name                = "%s"
  platform            = "%s"
  platform_credential = "%s"
  platform_principal  = "%s"
}
`, name, platform.Name, platform.Credential, platform.Principal)
}

func testAccAwsSnsApplicationConfig_basicAttribute(name string, platform *testAccAwsSnsApplicationPlatform, attributeKey, attributeValue string) string {
	if platform.Principal == "" {
		return fmt.Sprintf(`
resource "aws_sns_application" "test" {
  name                = "%s"
  platform            = "%s"
  platform_credential = "%s"
  %s                  = "%s"
}
`, name, platform.Name, platform.Credential, attributeKey, attributeValue)
	}
	return fmt.Sprintf(`
resource "aws_sns_application" "test" {
  name                = "%s"
  platform            = "%s"
  platform_credential = "%s"
  platform_principal  = "%s"
  %s                  = "%s"
}
`, name, platform.Name, platform.Credential, platform.Principal, attributeKey, attributeValue)
}

func testAccAwsSnsApplicationConfig_iamRoleAttribute(name string, platform *testAccAwsSnsApplicationPlatform, attributeKey, iamRoleName string) string {
	return fmt.Sprintf(`
resource "aws_iam_role" "test" {
  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": {
    "Effect": "Allow",
    "Principal": {"Service": "sns.amazonaws.com"},
    "Action": "sts:AssumeRole"
  }
}
EOF

  name = "%s"
}

resource "aws_iam_role_policy_attachment" "test" {
  policy_arn = "arn:aws:iam::aws:policy/CloudWatchLogsFullAccess"
  role       = "${aws_iam_role.test.id}"
}

%s
`, iamRoleName, testAccAwsSnsApplicationConfig_basicAttribute(name, platform, attributeKey, "${aws_iam_role.test.arn}"))
}

func testAccAwsSnsApplicationConfig_snsTopicAttribute(name string, platform *testAccAwsSnsApplicationPlatform, attributeKey, snsTopicName string) string {
	return fmt.Sprintf(`
resource "aws_sns_topic" "test" {
  name = "%s"
}

%s
`, snsTopicName, testAccAwsSnsApplicationConfig_basicAttribute(name, platform, attributeKey, "${aws_sns_topic.test.arn}"))
}
<|endoftext|>


<commit_before>package main

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/http/httputil"
	"time"

	"github.com/apache/incubator-trafficcontrol/lib/go-log"
)

var Authenticated = true
var NoAuth = false
// Routes returns the routes, and a catchall route for when no route matches.
func Routes(d ServerData) ([]Route, http.Handler, error) {
	routes := []Route{
		//ASNs
		{1.2,http.MethodGet, `asns-wip(\.json)?$`, ASNsHandler(d.DB),ServersPrivLevel, Authenticated,nil},
		//CDNs
		{ 1.2, http.MethodGet, `cdns-wip(\.json)?$`, cdnsHandler(d.DB), CDNsPrivLevel, Authenticated, nil},
		{ 1.2, http.MethodGet, `cdns/{cdn}/configs/monitoring(\.json)?$`, monitoringHandler(d.DB), MonitoringPrivLevel, Authenticated, nil},
		//Divisions
		{ 1.2, http.MethodGet, `divisions-wip(\.json)?$`, divisionsHandler(d.DB), DivisionsPrivLevel, Authenticated, nil},
		//HwInfo
		{ 1.2, http.MethodGet, `hwinfo-wip(\.json)?$`, hwInfoHandler(d.DB), HWInfoPrivLevel, Authenticated, nil},
		//Parameters
		{ 1.2, http.MethodGet, `parameters-wip(\.json)?$`, parametersHandler(d.DB), ParametersPrivLevel, Authenticated, nil},
		//Regions
		{ 1.2, http.MethodGet, `regions-wip(\.json)?$`, regionsHandler(d.DB), ServersPrivLevel, Authenticated, nil},
		{ 1.2, http.MethodGet, "regions-wip/{id}$", regionsHandler(d.DB), ServersPrivLevel, Authenticated, nil},
		//Servers
		{ 1.2, http.MethodGet, `servers-wip(\.json)?$`, serversHandler(d.DB), ServersPrivLevel, Authenticated, nil},
		{ 1.2, http.MethodGet, "servers-wip/{id}$", serversHandler(d.DB), ServersPrivLevel, Authenticated, nil},
		{ 1.2, http.MethodPost, "servers/{server}/deliveryservices$", assignDeliveryServicesToServerHandler(d.DB), PrivLevelOperations, Authenticated, nil},
		//System
		{ 1.2, http.MethodGet, `system/info-wip(\.json)?$`, systemInfoHandler(d.DB), SystemInfoPrivLevel, Authenticated, nil},
	}
	return routes, rootHandler(d), nil
}

// RootHandler returns the / handler for the service, which reverse-proxies the old Perl Traffic Ops
func rootHandler(d ServerData) http.Handler {
	tr := &http.Transport{
		TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		DialContext: (&net.Dialer{
			Timeout:   time.Duration(d.Config.ProxyTimeout) * time.Second,
			KeepAlive: time.Duration(d.Config.ProxyKeepAlive) * time.Second,
		}).DialContext,
		TLSHandshakeTimeout:   time.Duration(d.Config.ProxyTLSTimeout) * time.Second,
		ResponseHeaderTimeout: time.Duration(d.Config.ProxyReadHeaderTimeout) * time.Second,
		//Other knobs we can turn: ExpectContinueTimeout,IdleConnTimeout
	}
	rp := httputil.NewSingleHostReverseProxy(d.URL)
	rp.Transport = tr

	rp.ErrorLog = log.Error //if we don't provide a logger to the reverse proxy it logs to stdout/err and is lost when ran by a script.
	log.Debugf("our reverseProxy: %++v\n", rp)
	log.Debugf("our reverseProxy's transport: %++v\n", tr)
	loggingProxyHandler := wrapAccessLog(d.Secrets[0], rp)
	return loggingProxyHandler
}
<commit_msg>fixed the PrivLevel variable references<commit_after>package main

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/http/httputil"
	"time"

	"github.com/apache/incubator-trafficcontrol/lib/go-log"
)

var Authenticated = true
var NoAuth = false

// Routes returns the routes, and a catchall route for when no route matches.
func Routes(d ServerData) ([]Route, http.Handler, error) {
	routes := []Route{
		//ASNs
		{1.2, http.MethodGet, `asns-wip(\.json)?$`, ASNsHandler(d.DB), ASNsPrivLevel, Authenticated, nil},
		//CDNs
		{1.2, http.MethodGet, `cdns-wip(\.json)?$`, cdnsHandler(d.DB), CDNsPrivLevel, Authenticated, nil},
		{1.2, http.MethodGet, `cdns/{cdn}/configs/monitoring(\.json)?$`, monitoringHandler(d.DB), MonitoringPrivLevel, Authenticated, nil},
		//Divisions
		{1.2, http.MethodGet, `divisions-wip(\.json)?$`, divisionsHandler(d.DB), DivisionsPrivLevel, Authenticated, nil},
		//HwInfo
		{1.2, http.MethodGet, `hwinfo-wip(\.json)?$`, hwInfoHandler(d.DB), HWInfoPrivLevel, Authenticated, nil},
		//Parameters
		{1.2, http.MethodGet, `parameters-wip(\.json)?$`, parametersHandler(d.DB), ParametersPrivLevel, Authenticated, nil},
		//Regions
		{1.2, http.MethodGet, `regions-wip(\.json)?$`, regionsHandler(d.DB), RegionsPrivLevel, Authenticated, nil},
		{1.2, http.MethodGet, "regions-wip/{id}$", regionsHandler(d.DB), RegionsPrivLevel, Authenticated, nil},
		//Servers
		{1.2, http.MethodGet, `servers-wip(\.json)?$`, serversHandler(d.DB), ServersPrivLevel, Authenticated, nil},
		{1.2, http.MethodGet, "servers-wip/{id}$", serversHandler(d.DB), ServersPrivLevel, Authenticated, nil},
		{1.2, http.MethodPost, "servers/{server}/deliveryservices$", assignDeliveryServicesToServerHandler(d.DB), PrivLevelOperations, Authenticated, nil},

		//Statuses
		{1.2, http.MethodGet, `statuses-wip(\.json)?$`, statusesHandler(d.DB), StatusesPrivLevel, Authenticated, nil},
		{1.2, http.MethodGet, "statuses-wip/{id}$", statusesHandler(d.DB), StatusesPrivLevel, Authenticated, nil},
		//System
		{1.2, http.MethodGet, `system/info-wip(\.json)?$`, systemInfoHandler(d.DB), SystemInfoPrivLevel, Authenticated, nil},
	}
	return routes, rootHandler(d), nil
}

// RootHandler returns the / handler for the service, which reverse-proxies the old Perl Traffic Ops
func rootHandler(d ServerData) http.Handler {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		DialContext: (&net.Dialer{
			Timeout:   time.Duration(d.Config.ProxyTimeout) * time.Second,
			KeepAlive: time.Duration(d.Config.ProxyKeepAlive) * time.Second,
		}).DialContext,
		TLSHandshakeTimeout:   time.Duration(d.Config.ProxyTLSTimeout) * time.Second,
		ResponseHeaderTimeout: time.Duration(d.Config.ProxyReadHeaderTimeout) * time.Second,
		//Other knobs we can turn: ExpectContinueTimeout,IdleConnTimeout
	}
	rp := httputil.NewSingleHostReverseProxy(d.URL)
	rp.Transport = tr

	rp.ErrorLog = log.Error //if we don't provide a logger to the reverse proxy it logs to stdout/err and is lost when ran by a script.
	log.Debugf("our reverseProxy: %++v\n", rp)
	log.Debugf("our reverseProxy's transport: %++v\n", tr)
	loggingProxyHandler := wrapAccessLog(d.Secrets[0], rp)
	return loggingProxyHandler
}
<|endoftext|>


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

package dataaccess

import (
	"fmt"
	"github.com/andreaskoch/allmark2/common/content"
	"github.com/andreaskoch/allmark2/common/route"
)

type ItemType int

func (itemType ItemType) String() string {
	switch itemType {

	case TypePhysical:
		return "physical"

	case TypeVirtual:
		return "virtual"

	case TypeFileCollection:
		return "filecollection"

	default:
		return "unknown"

	}

	panic("Unreachable")
}

const (
	TypePhysical ItemType = iota
	TypeVirtual
	TypeFileCollection
)

// An Item represents a single document in a repository.
type Item struct {
	*content.ContentProvider
	itemType   ItemType
	route      route.Route
	filesFunc  func() []*File
	childsFunc func() []*Item
}

func NewPhysicalItem(route route.Route, contentProvider *content.ContentProvider, files func() []*File, childs func() []*Item) (*Item, error) {
	return newItem(TypePhysical, route, contentProvider, files, childs)
}

func NewVirtualItem(route route.Route, contentProvider *content.ContentProvider, files func() []*File, childs func() []*Item) (*Item, error) {
	return newItem(TypeVirtual, route, contentProvider, files, childs)
}

func NewFileCollectionItem(route route.Route, contentProvider *content.ContentProvider, files func() []*File) (*Item, error) {
	return newItem(TypeFileCollection, route, contentProvider, files, nil)
}

func newItem(itemType ItemType, route route.Route, contentProvider *content.ContentProvider, files func() []*File, childs func() []*Item) (*Item, error) {
	return &Item{
		contentProvider,
		itemType,
		route,
		files,
		childs,
	}, nil
}

func (item *Item) String() string {
	return fmt.Sprintf("%s", item.route.String())
}

// Get the type of this item (e.g. "physical", "virtual", ...)
func (item *Item) Type() ItemType {
	return item.itemType
}

// Gets a flag inidicating whether this item can have childs or not.
func (item *Item) CanHaveChilds() bool {
	switch item.Type() {

	// each child directory which is not the "files" folder can be a child
	case TypePhysical, TypeVirtual:
		return true

		// file collection items cannot have childs because all items in the directory are "files" and not items
	case TypeFileCollection:
		return false

	}

	panic("Unreachable. Unknown Item type.")
}

// Get the route of this item.
func (item *Item) Route() route.Route {
	return item.route
}

// Get the childs of this item. Returns nil if this item cannot have childs; otherwise returns a slice with zero or more childs.
func (item *Item) GetChilds() (childs []*Item) {
	if !item.CanHaveChilds() || item.childsFunc == nil {
		return
	}

	return item.childsFunc()
}

// Get the files of this item. Returns a slice of zero or more files.
func (item *Item) Files() (files []*File) {

	if item.filesFunc == nil {
		return []*File{}
	}

	return item.filesFunc()
}

func (item *Item) GetChildItemChanges(list []*Item) (newChilds []*Item, removedChilds []route.Route) {

	// capture the status quo
	previousChilds := make(map[string]*Item, 0)
	for _, child := range list {
		previousChilds[child.Route().Value()] = child
	}

	// get the new childs
	currentChilds := make(map[string]*Item, 0)
	for _, child := range item.GetChilds() {
		currentChilds[child.Route().Value()] = child
	}

	// find new childs
	newChilds = make([]*Item, 0)
	for key, child := range currentChilds {

		if _, exists := previousChilds[key]; !exists {
			newChilds = append(newChilds, child)
		}

	}

	// find removed childs
	removedChilds = make([]route.Route, 0)
	for key, child := range previousChilds {

		if _, exists := currentChilds[key]; !exists {
			removedChilds = append(removedChilds, child.Route())
		}

	}

	return newChilds, removedChilds

}
<commit_msg>Removed an unused method from the data access item<commit_after>// Copyright 2014 Andreas Koch. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package dataaccess

import (
	"fmt"
	"github.com/andreaskoch/allmark2/common/content"
	"github.com/andreaskoch/allmark2/common/route"
)

type ItemType int

func (itemType ItemType) String() string {
	switch itemType {

	case TypePhysical:
		return "physical"

	case TypeVirtual:
		return "virtual"

	case TypeFileCollection:
		return "filecollection"

	default:
		return "unknown"

	}

	panic("Unreachable")
}

const (
	TypePhysical ItemType = iota
	TypeVirtual
	TypeFileCollection
)

// An Item represents a single document in a repository.
type Item struct {
	*content.ContentProvider
	itemType   ItemType
	route      route.Route
	filesFunc  func() []*File
	childsFunc func() []*Item
}

func NewPhysicalItem(route route.Route, contentProvider *content.ContentProvider, files func() []*File, childs func() []*Item) (*Item, error) {
	return newItem(TypePhysical, route, contentProvider, files, childs)
}

func NewVirtualItem(route route.Route, contentProvider *content.ContentProvider, files func() []*File, childs func() []*Item) (*Item, error) {
	return newItem(TypeVirtual, route, contentProvider, files, childs)
}

func NewFileCollectionItem(route route.Route, contentProvider *content.ContentProvider, files func() []*File) (*Item, error) {
	return newItem(TypeFileCollection, route, contentProvider, files, nil)
}

func newItem(itemType ItemType, route route.Route, contentProvider *content.ContentProvider, files func() []*File, childs func() []*Item) (*Item, error) {
	return &Item{
		contentProvider,
		itemType,
		route,
		files,
		childs,
	}, nil
}

func (item *Item) String() string {
	return fmt.Sprintf("%s", item.route.String())
}

// Get the type of this item (e.g. "physical", "virtual", ...)
func (item *Item) Type() ItemType {
	return item.itemType
}

// Gets a flag inidicating whether this item can have childs or not.
func (item *Item) CanHaveChilds() bool {
	switch item.Type() {

	// each child directory which is not the "files" folder can be a child
	case TypePhysical, TypeVirtual:
		return true

		// file collection items cannot have childs because all items in the directory are "files" and not items
	case TypeFileCollection:
		return false

	}

	panic("Unreachable. Unknown Item type.")
}

// Get the route of this item.
func (item *Item) Route() route.Route {
	return item.route
}

// Get the childs of this item. Returns nil if this item cannot have childs; otherwise returns a slice with zero or more childs.
func (item *Item) GetChilds() (childs []*Item) {
	if !item.CanHaveChilds() || item.childsFunc == nil {
		return
	}

	return item.childsFunc()
}

// Get the files of this item. Returns a slice of zero or more files.
func (item *Item) Files() (files []*File) {

	if item.filesFunc == nil {
		return []*File{}
	}

	return item.filesFunc()
}
<|endoftext|>


<commit_before>// Copyright Project Harbor 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 period

import (
	"errors"
	"fmt"
	"math/rand"
	"time"

	"github.com/gocraft/work"
	"github.com/goharbor/harbor/src/jobservice/job"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/goharbor/harbor/src/jobservice/models"
	"github.com/goharbor/harbor/src/jobservice/opm"
	"github.com/goharbor/harbor/src/jobservice/utils"
	"github.com/gomodule/redigo/redis"
	"github.com/robfig/cron"
)

const (
	periodicEnqueuerSleep   = 2 * time.Minute
	periodicEnqueuerHorizon = 4 * time.Minute
)

type periodicEnqueuer struct {
	namespace        string
	pool             *redis.Pool
	policyStore      *periodicJobPolicyStore
	stopChan         chan struct{}
	doneStoppingChan chan struct{}
	statsManager     opm.JobStatsManager
}

func newPeriodicEnqueuer(namespace string, pool *redis.Pool, policyStore *periodicJobPolicyStore, statsManager opm.JobStatsManager) *periodicEnqueuer {
	return &periodicEnqueuer{
		namespace:        namespace,
		pool:             pool,
		policyStore:      policyStore,
		statsManager:     statsManager,
		stopChan:         make(chan struct{}),
		doneStoppingChan: make(chan struct{}),
	}
}

func (pe *periodicEnqueuer) start() {
	go pe.loop()
	logger.Info("Periodic enqueuer is started")
}

func (pe *periodicEnqueuer) stop() {
	pe.stopChan <- struct{}{}
	<-pe.doneStoppingChan
}

func (pe *periodicEnqueuer) loop() {
	defer func() {
		logger.Info("Periodic enqueuer is stopped")
	}()
	// Begin reaping periodically
	timer := time.NewTimer(periodicEnqueuerSleep + time.Duration(rand.Intn(30))*time.Second)
	defer timer.Stop()

	if pe.shouldEnqueue() {
		err := pe.enqueue()
		if err != nil {
			logger.Errorf("periodic_enqueuer.loop.enqueue:%s\n", err)
		}
	} else {
		logger.Debug("Enqueue condition not matched, do nothing.")
	}

	for {
		select {
		case <-pe.stopChan:
			pe.doneStoppingChan <- struct{}{}
			return
		case <-timer.C:
			timer.Reset(periodicEnqueuerSleep + time.Duration(rand.Intn(30))*time.Second)
			if pe.shouldEnqueue() {
				err := pe.enqueue()
				if err != nil {
					logger.Errorf("periodic_enqueuer.loop.enqueue:%s\n", err)
				}
			} else {
				logger.Debug("Enqueue condition not matched, do nothing.")
			}
		}
	}
}

func (pe *periodicEnqueuer) enqueue() error {
	now := time.Now().Unix()

	logger.Debugf("Periodic enqueuing loop: %d", now)

	conn := pe.pool.Get()
	defer conn.Close()

	// Set last periodic enqueue timestamp in advance to avoid duplicated enqueue actions
	if _, err := conn.Do("SET", utils.RedisKeyLastPeriodicEnqueue(pe.namespace), now); err != nil {
		return err
	}

	// Avoid schedule in the same time.
	lockerKey := fmt.Sprintf("%s:%s", utils.KeyPeriod(pe.namespace), "lock")
	lockerID := utils.MakeIdentifier()

	// Acquire a locker with 30s expiring time
	if err := acquireLock(conn, lockerKey, lockerID, 30); err != nil {
		return err
	}
	defer func() {
		if err := releaseLock(conn, lockerKey, lockerID); err != nil {
			logger.Errorf("release enqueue locker failed: %s", err)
		}
	}()

	nowTime := time.Unix(now, 0)
	horizon := nowTime.Add(periodicEnqueuerHorizon)

	for _, pl := range pe.policyStore.list() {
		schedule, err := cron.Parse(pl.CronSpec)
		if err != nil {
			// The cron spec should be already checked at top components.
			// Just in cases, if error occurred, ignore it
			logger.Errorf("[Ignore] Invalid corn spec in periodic policy %s %s: %s", pl.JobName, pl.PolicyID, err)
			continue
		}

		executions := []string{}
		for t := schedule.Next(nowTime); t.Before(horizon); t = schedule.Next(t) {
			epoch := t.Unix()
			scheduledExecutionID := utils.MakeIdentifier()
			executions = append(executions, scheduledExecutionID)

			// Create an execution (job) based on the periodic job template (policy)
			job := &work.Job{
				Name: pl.JobName,
				ID:   scheduledExecutionID,

				// This is technically wrong, but this lets the bytes be identical for the same periodic job instance.
				// If we don't do this, we'd need to use a different approach -- probably giving each periodic job its own
				// history of the past 100 periodic jobs, and only scheduling a job if it's not in the history.
				EnqueuedAt: epoch,
				Args:       pl.JobParameters, // Pass parameters to scheduled job here
			}

			rawJSON, err := utils.SerializeJob(job)
			if err != nil {
				return err
			}

			// Place the time slots for the job (policy)
			// If the slot is already there, error will be returned.
			expireTime := (epoch - nowTime.Unix()) + 5
			slot := fmt.Sprintf("%s:%s@%d", utils.KeyPeriodicJobTimeSlots(pe.namespace), pl.PolicyID, epoch)
			if err := placeSlot(conn, slot, epoch, expireTime); err != nil {
				// Logged and continue
				logger.Errorf("Failed to place time slot '%s@%d': %s", pl.PolicyID, epoch, err)
				continue
			}

			_, err = conn.Do("ZADD", utils.RedisKeyScheduled(pe.namespace), epoch, rawJSON)
			if err != nil {
				return err
			}

			logger.Infof("Schedule job %s:%s for policy %s at %d\n", job.Name, job.ID, pl.PolicyID, epoch)

			// Try to save the stats of new scheduled execution (job).
			pe.createExecution(pl.PolicyID, pl.JobName, scheduledExecutionID, epoch)

			// Get web hook from the periodic job (policy)
			webHookURL, err := pe.statsManager.GetHook(pl.PolicyID)
			if err == nil {
				// Register hook for the execution
				if err := pe.statsManager.RegisterHook(scheduledExecutionID, webHookURL, false); err != nil {
					// Just logged
					logger.Errorf("Failed to register web hook '%s' for periodic job (execution) '%s' with error: %s", webHookURL, scheduledExecutionID, err)
				}
			} else {
				// Just a warning
				logger.Warningf("Failed to retrieve web hook for periodic job (policy) %s: %s", pl.PolicyID, err)
			}
		}
		// Link the upstream job (policy) with the created executions
		if len(executions) > 0 {
			if err := pe.statsManager.AttachExecution(pl.PolicyID, executions...); err != nil {
				// Just logged it
				logger.Errorf("Link upstream job with executions failed: %s", err)
			}
		}
		// Directly use redis conn to update the periodic job (policy) status
		// Do not care the result
		conn.Do("HMSET", utils.KeyJobStats(pe.namespace, pl.PolicyID), "status", job.JobStatusScheduled, "update_time", time.Now().Unix())
	}

	return nil
}

func (pe *periodicEnqueuer) createExecution(upstreamJobID, upstreamJobName, executionID string, runAt int64) {
	execution := models.JobStats{
		Stats: &models.JobStatData{
			JobID:         executionID,
			JobName:       upstreamJobName,
			Status:        job.JobStatusPending,
			JobKind:       job.JobKindScheduled,
			EnqueueTime:   time.Now().Unix(),
			UpdateTime:    time.Now().Unix(),
			RefLink:       fmt.Sprintf("/api/v1/jobs/%s", executionID),
			RunAt:         runAt,
			UpstreamJobID: upstreamJobID,
		},
	}

	pe.statsManager.Save(execution)
}

func (pe *periodicEnqueuer) shouldEnqueue() bool {
	conn := pe.pool.Get()
	defer conn.Close()

	lastEnqueue, err := redis.Int64(conn.Do("GET", utils.RedisKeyLastPeriodicEnqueue(pe.namespace)))
	if err == redis.ErrNil {
		return true
	} else if err != nil {
		logger.Errorf("periodic_enqueuer.should_enqueue:%s\n", err)
		return true
	}

	return lastEnqueue < (time.Now().Unix() - int64(periodicEnqueuerSleep/time.Minute))
}

func placeSlot(conn redis.Conn, key string, value interface{}, expireTime int64) error {
	args := []interface{}{key, value, "NX", "EX", expireTime}
	res, err := conn.Do("SET", args...)
	if err != nil {
		return err
	}
	// Existing, the value can not be overrid
	if res == nil {
		return fmt.Errorf("key %s is already set with value %v", key, value)
	}

	return nil
}

func acquireLock(conn redis.Conn, lockerKey string, lockerID string, expireTime int64) error {
	return placeSlot(conn, lockerKey, lockerID, expireTime)
}

func releaseLock(conn redis.Conn, lockerKey string, lockerID string) error {
	theID, err := redis.String(conn.Do("GET", lockerKey))
	if err != nil {
		return err
	}

	if theID == lockerID {
		_, err := conn.Do("DEL", lockerKey)
		return err
	}

	return errors.New("locker ID mismatch")
}
<commit_msg>Use separate redis conn for the lock mechanism in the periodic enqueuer<commit_after>// Copyright Project Harbor 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 period

import (
	"errors"
	"fmt"
	"math/rand"
	"time"

	"github.com/gocraft/work"
	"github.com/goharbor/harbor/src/jobservice/job"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/goharbor/harbor/src/jobservice/models"
	"github.com/goharbor/harbor/src/jobservice/opm"
	"github.com/goharbor/harbor/src/jobservice/utils"
	"github.com/gomodule/redigo/redis"
	"github.com/robfig/cron"
)

const (
	periodicEnqueuerSleep   = 2 * time.Minute
	periodicEnqueuerHorizon = 4 * time.Minute
)

type periodicEnqueuer struct {
	namespace        string
	pool             *redis.Pool
	policyStore      *periodicJobPolicyStore
	stopChan         chan struct{}
	doneStoppingChan chan struct{}
	statsManager     opm.JobStatsManager
	identity         string
}

func newPeriodicEnqueuer(namespace string, pool *redis.Pool, policyStore *periodicJobPolicyStore, statsManager opm.JobStatsManager) *periodicEnqueuer {
	return &periodicEnqueuer{
		namespace:        namespace,
		pool:             pool,
		policyStore:      policyStore,
		statsManager:     statsManager,
		stopChan:         make(chan struct{}),
		doneStoppingChan: make(chan struct{}),
		identity:         utils.MakeIdentifier(), // Currently, use a generated ID
	}
}

func (pe *periodicEnqueuer) start() {
	go pe.loop()
	logger.Info("Periodic enqueuer is started")
}

func (pe *periodicEnqueuer) stop() {
	pe.stopChan <- struct{}{}
	<-pe.doneStoppingChan
}

func (pe *periodicEnqueuer) loop() {
	defer func() {
		logger.Info("Periodic enqueuer is stopped")
	}()
	// Begin reaping periodically
	timer := time.NewTimer(periodicEnqueuerSleep + time.Duration(rand.Intn(30))*time.Second)
	defer timer.Stop()

	if pe.shouldEnqueue() {
		err := pe.enqueue()
		if err != nil {
			logger.Errorf("periodic_enqueuer.loop.enqueue:%s\n", err)
		}
	} else {
		logger.Debug("Enqueue condition not matched, do nothing.")
	}

	for {
		select {
		case <-pe.stopChan:
			pe.doneStoppingChan <- struct{}{}
			return
		case <-timer.C:
			timer.Reset(periodicEnqueuerSleep + time.Duration(rand.Intn(30))*time.Second)
			if pe.shouldEnqueue() {
				err := pe.enqueue()
				if err != nil {
					logger.Errorf("periodic_enqueuer.loop.enqueue:%s\n", err)
				}
			} else {
				logger.Debug("Enqueue condition not matched, do nothing.")
			}
		}
	}
}

func (pe *periodicEnqueuer) enqueue() error {
	now := time.Now().Unix()

	logger.Debugf("Periodic enqueuing loop by enqueuer %s: %d", pe.identity, now)

	conn := pe.pool.Get()
	defer conn.Close()

	// Set last periodic enqueue timestamp in advance to avoid duplicated enqueue actions
	if _, err := conn.Do("SET", utils.RedisKeyLastPeriodicEnqueue(pe.namespace), now); err != nil {
		return err
	}

	// Avoid schedule in the same time.
	lockKey := fmt.Sprintf("%s:%s", utils.KeyPeriod(pe.namespace), "lock")

	// Use separate conn for the locker
	lockConn := pe.pool.Get()
	defer lockConn.Close()

	// Acquire a locker with 30s expiring time
	if err := acquireLock(lockConn, lockKey, pe.identity, 30); err != nil {
		return err
	}
	logger.Debugf("Periodic enqueuer %s acquires lock", pe.identity)

	defer func() {
		if err := releaseLock(lockConn, lockKey, pe.identity); err != nil {
			logger.Errorf("Periodic enqueuer %s releases lock failed: %s", pe.identity, err)
		} else {
			logger.Debugf("Periodic enqueuer %s releases lock", pe.identity)
		}
	}()

	nowTime := time.Unix(now, 0)
	horizon := nowTime.Add(periodicEnqueuerHorizon)

	for _, pl := range pe.policyStore.list() {
		schedule, err := cron.Parse(pl.CronSpec)
		if err != nil {
			// The cron spec should be already checked at top components.
			// Just in cases, if error occurred, ignore it
			logger.Errorf("[Ignore] Invalid corn spec in periodic policy %s %s: %s", pl.JobName, pl.PolicyID, err)
			continue
		}

		executions := []string{}
		for t := schedule.Next(nowTime); t.Before(horizon); t = schedule.Next(t) {
			epoch := t.Unix()
			scheduledExecutionID := utils.MakeIdentifier()
			executions = append(executions, scheduledExecutionID)

			// Create an execution (job) based on the periodic job template (policy)
			job := &work.Job{
				Name: pl.JobName,
				ID:   scheduledExecutionID,

				// This is technically wrong, but this lets the bytes be identical for the same periodic job instance.
				// If we don't do this, we'd need to use a different approach -- probably giving each periodic job its own
				// history of the past 100 periodic jobs, and only scheduling a job if it's not in the history.
				EnqueuedAt: epoch,
				Args:       pl.JobParameters, // Pass parameters to scheduled job here
			}

			rawJSON, err := utils.SerializeJob(job)
			if err != nil {
				return err
			}

			// Place the time slots for the job (policy)
			// If the slot is already there, error will be returned.
			expireTime := (epoch - nowTime.Unix()) + 5
			slot := fmt.Sprintf("%s:%s@%d", utils.KeyPeriodicJobTimeSlots(pe.namespace), pl.PolicyID, epoch)
			if err := placeSlot(conn, slot, epoch, expireTime); err != nil {
				// Logged and continue
				logger.Errorf("Failed to place time slot '%s@%d' in enqueuer %s: %s", pl.PolicyID, epoch, pe.identity, err)
				continue
			}

			_, err = conn.Do("ZADD", utils.RedisKeyScheduled(pe.namespace), epoch, rawJSON)
			if err != nil {
				return err
			}

			logger.Infof("Schedule job %s:%s for policy %s at %d by enqueuer %s", job.Name, job.ID, pl.PolicyID, epoch, pe.identity)

			// Try to save the stats of new scheduled execution (job).
			pe.createExecution(pl.PolicyID, pl.JobName, scheduledExecutionID, epoch)

			// Get web hook from the periodic job (policy)
			webHookURL, err := pe.statsManager.GetHook(pl.PolicyID)
			if err == nil {
				// Register hook for the execution
				if err := pe.statsManager.RegisterHook(scheduledExecutionID, webHookURL, false); err != nil {
					// Just logged
					logger.Errorf("Failed to register web hook '%s' for periodic job (execution) '%s' with error by enqueuer %s: %s", webHookURL, scheduledExecutionID, pe.identity, err)
				}
			} else {
				// Just a warning
				logger.Warningf("Failed to retrieve web hook for periodic job (policy) %s by enqueuer %s: %s", pl.PolicyID, pe.identity, err)
			}
		}
		// Link the upstream job (policy) with the created executions
		if len(executions) > 0 {
			if err := pe.statsManager.AttachExecution(pl.PolicyID, executions...); err != nil {
				// Just logged it
				logger.Errorf("Link upstream job with executions failed in enqueuer %s: %s", pe.identity, err)
			}
		}
		// Directly use redis conn to update the periodic job (policy) status
		// Do not care the result
		conn.Do("HMSET", utils.KeyJobStats(pe.namespace, pl.PolicyID), "status", job.JobStatusScheduled, "update_time", time.Now().Unix())
	}

	return nil
}

func (pe *periodicEnqueuer) createExecution(upstreamJobID, upstreamJobName, executionID string, runAt int64) {
	execution := models.JobStats{
		Stats: &models.JobStatData{
			JobID:         executionID,
			JobName:       upstreamJobName,
			Status:        job.JobStatusPending,
			JobKind:       job.JobKindScheduled,
			EnqueueTime:   time.Now().Unix(),
			UpdateTime:    time.Now().Unix(),
			RefLink:       fmt.Sprintf("/api/v1/jobs/%s", executionID),
			RunAt:         runAt,
			UpstreamJobID: upstreamJobID,
		},
	}

	pe.statsManager.Save(execution)
}

func (pe *periodicEnqueuer) shouldEnqueue() bool {
	conn := pe.pool.Get()
	defer conn.Close()

	lastEnqueue, err := redis.Int64(conn.Do("GET", utils.RedisKeyLastPeriodicEnqueue(pe.namespace)))
	if err == redis.ErrNil {
		return true
	} else if err != nil {
		logger.Errorf("periodic_enqueuer.should_enqueue:%s\n", err)
		return true
	}

	return lastEnqueue < (time.Now().Unix() - int64(periodicEnqueuerSleep/time.Minute))
}

func placeSlot(conn redis.Conn, key string, value interface{}, expireTime int64) error {
	args := []interface{}{key, value, "NX", "EX", expireTime}
	res, err := conn.Do("SET", args...)
	if err != nil {
		return err
	}
	// Existing, the value can not be overrid
	if res == nil {
		return fmt.Errorf("key %s is already set with value %v", key, value)
	}

	return nil
}

func acquireLock(conn redis.Conn, lockerKey string, lockerID string, expireTime int64) error {
	return placeSlot(conn, lockerKey, lockerID, expireTime)
}

func releaseLock(conn redis.Conn, lockerKey string, lockerID string) error {
	theID, err := redis.String(conn.Do("GET", lockerKey))
	if err != nil {
		return err
	}

	if theID == lockerID {
		_, err := conn.Do("DEL", lockerKey)
		return err
	}

	return errors.New("locker ID mismatch")
}
<|endoftext|>


<commit_before><commit_msg>No error for unhandled oobm<commit_after><|endoftext|>


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

// +build !windows,!nacl,!plan9,!js

package syslog

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"runtime"
	"sync"
	"testing"
	"time"
)

func runPktSyslog(c net.PacketConn, done chan<- string) {
	var buf [4096]byte
	var rcvd string
	ct := 0
	for {
		var n int
		var err error

		c.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
		n, _, err = c.ReadFrom(buf[:])
		rcvd += string(buf[:n])
		if err != nil {
			if oe, ok := err.(*net.OpError); ok {
				if ct < 3 && oe.Temporary() {
					ct++
					continue
				}
			}
			break
		}
	}
	c.Close()
	done <- rcvd
}

var crashy = false

func testableNetwork(network string) bool {
	switch network {
	case "unix", "unixgram":
		switch runtime.GOOS {
		case "darwin":
			switch runtime.GOARCH {
			case "arm", "arm64":
				return false
			}
		case "android":
			return false
		}
	}
	return true
}

func runStreamSyslog(l net.Listener, done chan<- string, wg *sync.WaitGroup) {
	for {
		var c net.Conn
		var err error
		if c, err = l.Accept(); err != nil {
			return
		}
		wg.Add(1)
		go func(c net.Conn) {
			defer wg.Done()
			c.SetReadDeadline(time.Now().Add(5 * time.Second))
			b := bufio.NewReader(c)
			for ct := 1; !crashy || ct&7 != 0; ct++ {
				s, err := b.ReadString('\n')
				if err != nil {
					break
				}
				done <- s
			}
			c.Close()
		}(c)
	}
}

func startServer(n, la string, done chan<- string) (addr string, sock io.Closer, wg *sync.WaitGroup) {
	if n == "udp" || n == "tcp" {
		la = "127.0.0.1:0"
	} else {
		// unix and unixgram: choose an address if none given
		if la == "" {
			// use ioutil.TempFile to get a name that is unique
			f, err := ioutil.TempFile("", "syslogtest")
			if err != nil {
				log.Fatal("TempFile: ", err)
			}
			f.Close()
			la = f.Name()
		}
		os.Remove(la)
	}

	wg = new(sync.WaitGroup)
	if n == "udp" || n == "unixgram" {
		l, e := net.ListenPacket(n, la)
		if e != nil {
			log.Fatalf("startServer failed: %v", e)
		}
		addr = l.LocalAddr().String()
		sock = l
		wg.Add(1)
		go func() {
			defer wg.Done()
			runPktSyslog(l, done)
		}()
	} else {
		l, e := net.Listen(n, la)
		if e != nil {
			log.Fatalf("startServer failed: %v", e)
		}
		addr = l.Addr().String()
		sock = l
		wg.Add(1)
		go func() {
			defer wg.Done()
			runStreamSyslog(l, done, wg)
		}()
	}
	return
}

func TestWithSimulated(t *testing.T) {
	t.Parallel()
	msg := "Test 123"
	var transport []string
	for _, n := range []string{"unix", "unixgram", "udp", "tcp"} {
		if testableNetwork(n) {
			transport = append(transport, n)
		}
	}

	for _, tr := range transport {
		done := make(chan string)
		addr, sock, srvWG := startServer(tr, "", done)
		defer srvWG.Wait()
		defer sock.Close()
		if tr == "unix" || tr == "unixgram" {
			defer os.Remove(addr)
		}
		s, err := Dial(tr, addr, LOG_INFO|LOG_USER, "syslog_test")
		if err != nil {
			t.Fatalf("Dial() failed: %v", err)
		}
		err = s.Info(msg)
		if err != nil {
			t.Fatalf("log failed: %v", err)
		}
		check(t, msg, <-done)
		s.Close()
	}
}

func TestFlap(t *testing.T) {
	net := "unix"
	if !testableNetwork(net) {
		t.Skipf("skipping on %s/%s; 'unix' is not supported", runtime.GOOS, runtime.GOARCH)
	}

	done := make(chan string)
	addr, sock, srvWG := startServer(net, "", done)
	defer srvWG.Wait()
	defer os.Remove(addr)
	defer sock.Close()

	s, err := Dial(net, addr, LOG_INFO|LOG_USER, "syslog_test")
	if err != nil {
		t.Fatalf("Dial() failed: %v", err)
	}
	msg := "Moo 2"
	err = s.Info(msg)
	if err != nil {
		t.Fatalf("log failed: %v", err)
	}
	check(t, msg, <-done)

	// restart the server
	_, sock2, srvWG2 := startServer(net, addr, done)
	defer srvWG2.Wait()
	defer sock2.Close()

	// and try retransmitting
	msg = "Moo 3"
	err = s.Info(msg)
	if err != nil {
		t.Fatalf("log failed: %v", err)
	}
	check(t, msg, <-done)

	s.Close()
}

func TestNew(t *testing.T) {
	if LOG_LOCAL7 != 23<<3 {
		t.Fatalf("LOG_LOCAL7 has wrong value")
	}
	if testing.Short() {
		// Depends on syslog daemon running, and sometimes it's not.
		t.Skip("skipping syslog test during -short")
	}

	s, err := New(LOG_INFO|LOG_USER, "the_tag")
	if err != nil {
		t.Fatalf("New() failed: %s", err)
	}
	// Don't send any messages.
	s.Close()
}

func TestNewLogger(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping syslog test during -short")
	}
	f, err := NewLogger(LOG_USER|LOG_INFO, 0)
	if f == nil {
		t.Error(err)
	}
}

func TestDial(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping syslog test during -short")
	}
	f, err := Dial("", "", (LOG_LOCAL7|LOG_DEBUG)+1, "syslog_test")
	if f != nil {
		t.Fatalf("Should have trapped bad priority")
	}
	f, err = Dial("", "", -1, "syslog_test")
	if f != nil {
		t.Fatalf("Should have trapped bad priority")
	}
	l, err := Dial("", "", LOG_USER|LOG_ERR, "syslog_test")
	if err != nil {
		t.Fatalf("Dial() failed: %s", err)
	}
	l.Close()
}

func check(t *testing.T, in, out string) {
	tmpl := fmt.Sprintf("<%d>%%s %%s syslog_test[%%d]: %s\n", LOG_USER+LOG_INFO, in)
	if hostname, err := os.Hostname(); err != nil {
		t.Error("Error retrieving hostname")
	} else {
		var parsedHostname, timestamp string
		var pid int
		if n, err := fmt.Sscanf(out, tmpl, &timestamp, &parsedHostname, &pid); n != 3 || err != nil || hostname != parsedHostname {
			t.Errorf("Got %q, does not match template %q (%d %s)", out, tmpl, n, err)
		}
	}
}

func TestWrite(t *testing.T) {
	t.Parallel()
	tests := []struct {
		pri Priority
		pre string
		msg string
		exp string
	}{
		{LOG_USER | LOG_ERR, "syslog_test", "", "%s %s syslog_test[%d]: \n"},
		{LOG_USER | LOG_ERR, "syslog_test", "write test", "%s %s syslog_test[%d]: write test\n"},
		// Write should not add \n if there already is one
		{LOG_USER | LOG_ERR, "syslog_test", "write test 2\n", "%s %s syslog_test[%d]: write test 2\n"},
	}

	if hostname, err := os.Hostname(); err != nil {
		t.Fatalf("Error retrieving hostname")
	} else {
		for _, test := range tests {
			done := make(chan string)
			addr, sock, srvWG := startServer("udp", "", done)
			defer srvWG.Wait()
			defer sock.Close()
			l, err := Dial("udp", addr, test.pri, test.pre)
			if err != nil {
				t.Fatalf("syslog.Dial() failed: %v", err)
			}
			defer l.Close()
			_, err = io.WriteString(l, test.msg)
			if err != nil {
				t.Fatalf("WriteString() failed: %v", err)
			}
			rcvd := <-done
			test.exp = fmt.Sprintf("<%d>", test.pri) + test.exp
			var parsedHostname, timestamp string
			var pid int
			if n, err := fmt.Sscanf(rcvd, test.exp, &timestamp, &parsedHostname, &pid); n != 3 || err != nil || hostname != parsedHostname {
				t.Errorf("s.Info() = '%q', didn't match '%q' (%d %s)", rcvd, test.exp, n, err)
			}
		}
	}
}

func TestConcurrentWrite(t *testing.T) {
	addr, sock, srvWG := startServer("udp", "", make(chan string, 1))
	defer srvWG.Wait()
	defer sock.Close()
	w, err := Dial("udp", addr, LOG_USER|LOG_ERR, "how's it going?")
	if err != nil {
		t.Fatalf("syslog.Dial() failed: %v", err)
	}
	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := w.Info("test")
			if err != nil {
				t.Errorf("Info() failed: %v", err)
				return
			}
		}()
	}
	wg.Wait()
}

func TestConcurrentReconnect(t *testing.T) {
	crashy = true
	defer func() { crashy = false }()

	const N = 10
	const M = 100
	net := "unix"
	if !testableNetwork(net) {
		net = "tcp"
		if !testableNetwork(net) {
			t.Skipf("skipping on %s/%s; neither 'unix' or 'tcp' is supported", runtime.GOOS, runtime.GOARCH)
		}
	}
	done := make(chan string, N*M)
	addr, sock, srvWG := startServer(net, "", done)
	if net == "unix" {
		defer os.Remove(addr)
	}

	// count all the messages arriving
	count := make(chan int)
	go func() {
		ct := 0
		for range done {
			ct++
			// we are looking for 500 out of 1000 events
			// here because lots of log messages are lost
			// in buffers (kernel and/or bufio)
			if ct > N*M/2 {
				break
			}
		}
		count <- ct
	}()

	var wg sync.WaitGroup
	wg.Add(N)
	for i := 0; i < N; i++ {
		go func() {
			defer wg.Done()
			w, err := Dial(net, addr, LOG_USER|LOG_ERR, "tag")
			if err != nil {
				t.Errorf("syslog.Dial() failed: %v", err)
				return
			}
			defer w.Close()
			for i := 0; i < M; i++ {
				err := w.Info("test")
				if err != nil {
					t.Errorf("Info() failed: %v", err)
					return
				}
			}
		}()
	}
	wg.Wait()
	sock.Close()
	srvWG.Wait()
	close(done)

	select {
	case <-count:
	case <-time.After(100 * time.Millisecond):
		t.Error("timeout in concurrent reconnect")
	}
}
<commit_msg>log/syslog: skip tests that depend on daemon on builders<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !windows,!nacl,!plan9,!js

package syslog

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"runtime"
	"sync"
	"testing"
	"time"
)

func runPktSyslog(c net.PacketConn, done chan<- string) {
	var buf [4096]byte
	var rcvd string
	ct := 0
	for {
		var n int
		var err error

		c.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
		n, _, err = c.ReadFrom(buf[:])
		rcvd += string(buf[:n])
		if err != nil {
			if oe, ok := err.(*net.OpError); ok {
				if ct < 3 && oe.Temporary() {
					ct++
					continue
				}
			}
			break
		}
	}
	c.Close()
	done <- rcvd
}

var crashy = false

func testableNetwork(network string) bool {
	switch network {
	case "unix", "unixgram":
		switch runtime.GOOS {
		case "darwin":
			switch runtime.GOARCH {
			case "arm", "arm64":
				return false
			}
		case "android":
			return false
		}
	}
	return true
}

func runStreamSyslog(l net.Listener, done chan<- string, wg *sync.WaitGroup) {
	for {
		var c net.Conn
		var err error
		if c, err = l.Accept(); err != nil {
			return
		}
		wg.Add(1)
		go func(c net.Conn) {
			defer wg.Done()
			c.SetReadDeadline(time.Now().Add(5 * time.Second))
			b := bufio.NewReader(c)
			for ct := 1; !crashy || ct&7 != 0; ct++ {
				s, err := b.ReadString('\n')
				if err != nil {
					break
				}
				done <- s
			}
			c.Close()
		}(c)
	}
}

func startServer(n, la string, done chan<- string) (addr string, sock io.Closer, wg *sync.WaitGroup) {
	if n == "udp" || n == "tcp" {
		la = "127.0.0.1:0"
	} else {
		// unix and unixgram: choose an address if none given
		if la == "" {
			// use ioutil.TempFile to get a name that is unique
			f, err := ioutil.TempFile("", "syslogtest")
			if err != nil {
				log.Fatal("TempFile: ", err)
			}
			f.Close()
			la = f.Name()
		}
		os.Remove(la)
	}

	wg = new(sync.WaitGroup)
	if n == "udp" || n == "unixgram" {
		l, e := net.ListenPacket(n, la)
		if e != nil {
			log.Fatalf("startServer failed: %v", e)
		}
		addr = l.LocalAddr().String()
		sock = l
		wg.Add(1)
		go func() {
			defer wg.Done()
			runPktSyslog(l, done)
		}()
	} else {
		l, e := net.Listen(n, la)
		if e != nil {
			log.Fatalf("startServer failed: %v", e)
		}
		addr = l.Addr().String()
		sock = l
		wg.Add(1)
		go func() {
			defer wg.Done()
			runStreamSyslog(l, done, wg)
		}()
	}
	return
}

func TestWithSimulated(t *testing.T) {
	t.Parallel()
	msg := "Test 123"
	var transport []string
	for _, n := range []string{"unix", "unixgram", "udp", "tcp"} {
		if testableNetwork(n) {
			transport = append(transport, n)
		}
	}

	for _, tr := range transport {
		done := make(chan string)
		addr, sock, srvWG := startServer(tr, "", done)
		defer srvWG.Wait()
		defer sock.Close()
		if tr == "unix" || tr == "unixgram" {
			defer os.Remove(addr)
		}
		s, err := Dial(tr, addr, LOG_INFO|LOG_USER, "syslog_test")
		if err != nil {
			t.Fatalf("Dial() failed: %v", err)
		}
		err = s.Info(msg)
		if err != nil {
			t.Fatalf("log failed: %v", err)
		}
		check(t, msg, <-done)
		s.Close()
	}
}

func TestFlap(t *testing.T) {
	net := "unix"
	if !testableNetwork(net) {
		t.Skipf("skipping on %s/%s; 'unix' is not supported", runtime.GOOS, runtime.GOARCH)
	}

	done := make(chan string)
	addr, sock, srvWG := startServer(net, "", done)
	defer srvWG.Wait()
	defer os.Remove(addr)
	defer sock.Close()

	s, err := Dial(net, addr, LOG_INFO|LOG_USER, "syslog_test")
	if err != nil {
		t.Fatalf("Dial() failed: %v", err)
	}
	msg := "Moo 2"
	err = s.Info(msg)
	if err != nil {
		t.Fatalf("log failed: %v", err)
	}
	check(t, msg, <-done)

	// restart the server
	_, sock2, srvWG2 := startServer(net, addr, done)
	defer srvWG2.Wait()
	defer sock2.Close()

	// and try retransmitting
	msg = "Moo 3"
	err = s.Info(msg)
	if err != nil {
		t.Fatalf("log failed: %v", err)
	}
	check(t, msg, <-done)

	s.Close()
}

func TestNew(t *testing.T) {
	if LOG_LOCAL7 != 23<<3 {
		t.Fatalf("LOG_LOCAL7 has wrong value")
	}
	if testing.Short() {
		// Depends on syslog daemon running, and sometimes it's not.
		t.Skip("skipping syslog test during -short")
	}

	s, err := New(LOG_INFO|LOG_USER, "the_tag")
	if err != nil {
		if err.Error() == "Unix syslog delivery error" {
			t.Skip("skipping: syslogd not running")
		}
		t.Fatalf("New() failed: %s", err)
	}
	// Don't send any messages.
	s.Close()
}

func TestNewLogger(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping syslog test during -short")
	}
	f, err := NewLogger(LOG_USER|LOG_INFO, 0)
	if f == nil {
		if err.Error() == "Unix syslog delivery error" {
			t.Skip("skipping: syslogd not running")
		}
		t.Error(err)
	}
}

func TestDial(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping syslog test during -short")
	}
	f, err := Dial("", "", (LOG_LOCAL7|LOG_DEBUG)+1, "syslog_test")
	if f != nil {
		t.Fatalf("Should have trapped bad priority")
	}
	f, err = Dial("", "", -1, "syslog_test")
	if f != nil {
		t.Fatalf("Should have trapped bad priority")
	}
	l, err := Dial("", "", LOG_USER|LOG_ERR, "syslog_test")
	if err != nil {
		if err.Error() == "Unix syslog delivery error" {
			t.Skip("skipping: syslogd not running")
		}
		t.Fatalf("Dial() failed: %s", err)
	}
	l.Close()
}

func check(t *testing.T, in, out string) {
	tmpl := fmt.Sprintf("<%d>%%s %%s syslog_test[%%d]: %s\n", LOG_USER+LOG_INFO, in)
	if hostname, err := os.Hostname(); err != nil {
		t.Error("Error retrieving hostname")
	} else {
		var parsedHostname, timestamp string
		var pid int
		if n, err := fmt.Sscanf(out, tmpl, &timestamp, &parsedHostname, &pid); n != 3 || err != nil || hostname != parsedHostname {
			t.Errorf("Got %q, does not match template %q (%d %s)", out, tmpl, n, err)
		}
	}
}

func TestWrite(t *testing.T) {
	t.Parallel()
	tests := []struct {
		pri Priority
		pre string
		msg string
		exp string
	}{
		{LOG_USER | LOG_ERR, "syslog_test", "", "%s %s syslog_test[%d]: \n"},
		{LOG_USER | LOG_ERR, "syslog_test", "write test", "%s %s syslog_test[%d]: write test\n"},
		// Write should not add \n if there already is one
		{LOG_USER | LOG_ERR, "syslog_test", "write test 2\n", "%s %s syslog_test[%d]: write test 2\n"},
	}

	if hostname, err := os.Hostname(); err != nil {
		t.Fatalf("Error retrieving hostname")
	} else {
		for _, test := range tests {
			done := make(chan string)
			addr, sock, srvWG := startServer("udp", "", done)
			defer srvWG.Wait()
			defer sock.Close()
			l, err := Dial("udp", addr, test.pri, test.pre)
			if err != nil {
				t.Fatalf("syslog.Dial() failed: %v", err)
			}
			defer l.Close()
			_, err = io.WriteString(l, test.msg)
			if err != nil {
				t.Fatalf("WriteString() failed: %v", err)
			}
			rcvd := <-done
			test.exp = fmt.Sprintf("<%d>", test.pri) + test.exp
			var parsedHostname, timestamp string
			var pid int
			if n, err := fmt.Sscanf(rcvd, test.exp, &timestamp, &parsedHostname, &pid); n != 3 || err != nil || hostname != parsedHostname {
				t.Errorf("s.Info() = '%q', didn't match '%q' (%d %s)", rcvd, test.exp, n, err)
			}
		}
	}
}

func TestConcurrentWrite(t *testing.T) {
	addr, sock, srvWG := startServer("udp", "", make(chan string, 1))
	defer srvWG.Wait()
	defer sock.Close()
	w, err := Dial("udp", addr, LOG_USER|LOG_ERR, "how's it going?")
	if err != nil {
		t.Fatalf("syslog.Dial() failed: %v", err)
	}
	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := w.Info("test")
			if err != nil {
				t.Errorf("Info() failed: %v", err)
				return
			}
		}()
	}
	wg.Wait()
}

func TestConcurrentReconnect(t *testing.T) {
	crashy = true
	defer func() { crashy = false }()

	const N = 10
	const M = 100
	net := "unix"
	if !testableNetwork(net) {
		net = "tcp"
		if !testableNetwork(net) {
			t.Skipf("skipping on %s/%s; neither 'unix' or 'tcp' is supported", runtime.GOOS, runtime.GOARCH)
		}
	}
	done := make(chan string, N*M)
	addr, sock, srvWG := startServer(net, "", done)
	if net == "unix" {
		defer os.Remove(addr)
	}

	// count all the messages arriving
	count := make(chan int)
	go func() {
		ct := 0
		for range done {
			ct++
			// we are looking for 500 out of 1000 events
			// here because lots of log messages are lost
			// in buffers (kernel and/or bufio)
			if ct > N*M/2 {
				break
			}
		}
		count <- ct
	}()

	var wg sync.WaitGroup
	wg.Add(N)
	for i := 0; i < N; i++ {
		go func() {
			defer wg.Done()
			w, err := Dial(net, addr, LOG_USER|LOG_ERR, "tag")
			if err != nil {
				t.Errorf("syslog.Dial() failed: %v", err)
				return
			}
			defer w.Close()
			for i := 0; i < M; i++ {
				err := w.Info("test")
				if err != nil {
					t.Errorf("Info() failed: %v", err)
					return
				}
			}
		}()
	}
	wg.Wait()
	sock.Close()
	srvWG.Wait()
	close(done)

	select {
	case <-count:
	case <-time.After(100 * time.Millisecond):
		t.Error("timeout in concurrent reconnect")
	}
}
<|endoftext|>


<commit_before>package gfm

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"net/http"
)

// Database is the interface of a FileMaker database
type Database interface {
	Server() Server
	Name() string
	Lay(string) Layout
	Lays() ([]Layout, error)
	// TODO: Find out if this is the proper way of representing script. Maybe
	// they need their own type?
	Scripts() ([]string, error)
	URL() string
}

type database struct {
	server Server
	name   string
}

func (db database) Server() Server {
	return db.server
}

func (db database) Name() string {
	return db.name
}

func (db database) Lay(name string) Layout {
	return layout{database: db, name: name}
}

func (db database) Lays() ([]Layout, error) {
	var lays []Layout

	var b bytes.Buffer
	b.WriteString(db.URL())
	b.WriteString("-layoutnames")
	res, err := http.Get(b.String())

	if err != nil {
		// TODO: maybe we need a better error here
		return lays, err
	}

	var fmrs FMResultSet
	err = xml.NewDecoder(res.Body).Decode(&fmrs)
	if err != nil {
		return lays, err
	} else if fmrs.Error.Code != 0 {
		return lays, fmt.Errorf("FMError: %v", fmrs.Error.Code)
	}

	for _, r := range fmrs.ResultSet.Records {
		name := r.Fields[0].Data
		lay := layout{database: db, name: name}
		lays = append(lays, lay)
	}

	return lays, nil
}

func (db database) Scripts() ([]string, error) {
	var scripts []string

	var b bytes.Buffer
	b.WriteString(db.URL())
	b.WriteString("-scriptnames")
	res, err := http.Get(b.String())

	if err != nil {
		// TODO: maybe we need a better error here
		return scripts, err
	}
	var fmrs FMResultSet
	err = xml.NewDecoder(res.Body).Decode(&fmrs)
	if err != nil {
		return scripts, err
	} else if fmrs.Error.Code != 0 {
		return scripts, fmt.Errorf("FMError: %v", fmrs.Error.Code)
	}

	for _, r := range fmrs.ResultSet.Records {
		name := r.Fields[0].Data
		scripts = append(scripts, name)
	}

	return scripts, nil
}

func (db database) URL() string {
	var b bytes.Buffer
	b.WriteString(db.Server().URL())
	b.WriteString("-db=")
	b.WriteString(db.Name())
	b.WriteString("&")
	return b.String()
}
<commit_msg>Document Database interface<commit_after>package gfm

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"net/http"
)

// Database is the interface of a FileMaker database
type Database interface {
	// Server returns the Server that this database belongs to.
	Server() Server

	// Name returns the name of the database.
	Name() string

	// Lay returns a Layout struct, which can be queried further.
	Lay(string) Layout

	// Lays returns the layouts of the database.
	Lays() ([]Layout, error)

	// Scripts returns an array of the scripts on this database.
	//
	// TODO: Find out if this is the proper way of representing script. Maybe
	// they need their own type?
	Scripts() ([]string, error)

	// URL returns the url for sending requests to this database.
	URL() string
}

type database struct {
	server Server
	name   string
}

func (db database) Server() Server {
	return db.server
}

func (db database) Name() string {
	return db.name
}

func (db database) Lay(name string) Layout {
	return layout{database: db, name: name}
}

func (db database) Lays() ([]Layout, error) {
	var lays []Layout

	var b bytes.Buffer
	b.WriteString(db.URL())
	b.WriteString("-layoutnames")
	res, err := http.Get(b.String())

	if err != nil {
		// TODO: maybe we need a better error here
		return lays, err
	}

	var fmrs FMResultSet
	err = xml.NewDecoder(res.Body).Decode(&fmrs)
	if err != nil {
		return lays, err
	} else if fmrs.Error.Code != 0 {
		return lays, fmt.Errorf("FMError: %v", fmrs.Error.Code)
	}

	for _, r := range fmrs.ResultSet.Records {
		name := r.Fields[0].Data
		lay := layout{database: db, name: name}
		lays = append(lays, lay)
	}

	return lays, nil
}

func (db database) Scripts() ([]string, error) {
	var scripts []string

	var b bytes.Buffer
	b.WriteString(db.URL())
	b.WriteString("-scriptnames")
	res, err := http.Get(b.String())

	if err != nil {
		// TODO: maybe we need a better error here
		return scripts, err
	}
	var fmrs FMResultSet
	err = xml.NewDecoder(res.Body).Decode(&fmrs)
	if err != nil {
		return scripts, err
	} else if fmrs.Error.Code != 0 {
		return scripts, fmt.Errorf("FMError: %v", fmrs.Error.Code)
	}

	for _, r := range fmrs.ResultSet.Records {
		name := r.Fields[0].Data
		scripts = append(scripts, name)
	}

	return scripts, nil
}

func (db database) URL() string {
	var b bytes.Buffer
	b.WriteString(db.Server().URL())
	b.WriteString("-db=")
	b.WriteString(db.Name())
	b.WriteString("&")
	return b.String()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"database/sql"
	"net/url"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/mattn/go-sqlite3"
	"go.uber.org/zap"

	"magneticod/bittorrent"

	"path"
	"os"
)

type engineType uint8

const (
	SQLITE     engineType = 0
	POSTGRESQL            = 1
	MYSQL                 = 2
)

type Database struct {
	database    *sql.DB
	engine      engineType
	newTorrents chan bittorrent.Metadata
}

// NewDatabase creates a new Database.
//
// url either starts with "sqlite:" or "postgresql:"
func NewDatabase(rawurl string) (*Database, error) {
	db := Database{}

	dbURL, err := url.Parse(rawurl)
	if err != nil {
		return nil, err
	}

	switch dbURL.Scheme {
	case "sqlite":
		db.engine = SQLITE
		// All this pain is to make sure that an empty file exist (in case the database is not there
		// yet) so that sql.Open won't fail.
		dbDir, _ := path.Split(dbURL.Path)
		if err := os.MkdirAll(dbDir, 0755); err != nil {
			return nil, fmt.Errorf("for directory `%s`:  %s", dbDir, err.Error())
		}
		f, err := os.OpenFile(dbURL.Path, os.O_CREATE, 0666)
		if err != nil {
			return nil, fmt.Errorf("for file `%s`: %s", dbURL.Path, err.Error())
		}
		if err := f.Sync(); err != nil {
			return nil, fmt.Errorf("for file `%s`: %s", dbURL.Path, err.Error())
		}
		if err := f.Close(); err != nil {
			return nil, fmt.Errorf("for file `%s`: %s", dbURL.Path, err.Error())
		}
		db.database, err = sql.Open("sqlite3", dbURL.RawPath)

	case "postgresql":
		db.engine = POSTGRESQL
		db.database, err = sql.Open("postgresql", rawurl)

	case "mysql":
		db.engine = MYSQL
		db.database, err = sql.Open("mysql", rawurl)

	default:
		return nil, fmt.Errorf("unknown URI scheme (or malformed URI)!")
	}

	// Check for errors from sql.Open()
	if err != nil {
		return nil, fmt.Errorf("sql.Open()!  %s", err.Error())
	}

	if err = db.database.Ping(); err != nil {
		return nil, fmt.Errorf("DB.Ping()!  %s", err.Error())
	}

	if err := db.setupDatabase(); err != nil {
		return nil, fmt.Errorf("setupDatabase()!  %s", err.Error())
	}

	db.newTorrents = make(chan bittorrent.Metadata, 10)

	return &db, nil
}

// AddNewTorrent adds a new torrent to the *queue* to be flushed to the persistent database.
func (db *Database) AddNewTorrent(torrent bittorrent.Metadata) error {
	for {
		select {
		case db.newTorrents <- torrent:
			return nil
		default:
			// newTorrents queue was full: flush and try again and again (and again)...
			err := db.flushNewTorrents()
			if err != nil {
				return err
			}
			continue
		}
	}
}

func (db *Database) flushNewTorrents() error {
	tx, err := db.database.Begin()
	if err != nil {
		return fmt.Errorf("sql.DB.Begin()!  %s", err.Error())
	}

	var nTorrents, nFiles uint
	for torrent := range db.newTorrents {
		res, err := tx.Exec("INSERT INTO torrents (info_hash, name, total_size, discovered_on) VALUES ($1, $2, $3, $4);",
			torrent.InfoHash, torrent.Name, torrent.TotalSize, torrent.DiscoveredOn)
		if err != nil {
			ourError := fmt.Errorf("error while INSERTing INTO torrents!  %s", err.Error())
			if err := tx.Rollback(); err != nil {
				return fmt.Errorf("%s\tmeanwhile, could not rollback the current transaction either!  %s", ourError.Error(), err.Error())
			}
			return ourError
		}
		var lastInsertId int64
		if lastInsertId, err = res.LastInsertId(); err != nil {
			return fmt.Errorf("sql.Result.LastInsertId()!  %s", err.Error())
		}

		for _, file := range torrent.Files {
			_, err := tx.Exec("INSERT INTO files (torrent_id, size, path) VALUES($1, $2, $3);",
			lastInsertId, file.Length, file.Path)
			if err != nil {
				ourError := fmt.Errorf("error while INSERTing INTO files!  %s", err.Error())
				if err := tx.Rollback(); err != nil {
					return fmt.Errorf("%s\tmeanwhile, could not rollback the current transaction either!  %s", ourError.Error(), err.Error())
				}
				return ourError
			}
			nFiles++
		}
		nTorrents++
	}

	err = tx.Commit()
	if err != nil {
		return fmt.Errorf("sql.Tx.Commit()!  %s", err.Error())
	}

	zap.L().Sugar().Infof("%d torrents (%d files) are flushed to the database successfully.",
		nTorrents, nFiles)
	return nil
}

func (db *Database) Close() {
	// Be careful to not to get into an infinite loop. =)
	db.database.Close()
}

func (db *Database) setupDatabase() error {
	switch db.engine {
	case SQLITE:
		return setupSqliteDatabase(db.database)

	case POSTGRESQL:
		zap.L().Fatal("setupDatabase() is not implemented for PostgreSQL yet!")

	case MYSQL:
		return setupMySQLDatabase(db.database)

	default:
		zap.L().Sugar().Fatalf("Unknown database engine value %d! (programmer error)", db.engine)
	}

	return nil
}

func setupSqliteDatabase(database *sql.DB) error {
	// Enable Write-Ahead Logging for SQLite as "WAL provides more concurrency as readers do not
	// block writers and a writer does not block readers. Reading and writing can proceed
	// concurrently."
	// Caveats:
	//   * Might be unsupported by OSes other than Windows and UNIXes.
	//   * Does not work over a network filesystem.
	//   * Transactions that involve changes against multiple ATTACHed databases are not atomic
	//     across all databases as a set.
	// See: https://www.sqlite.org/wal.html
	//
	// Force SQLite to use disk, instead of memory, for all temporary files to reduce the memory
	// footprint.
	//
	// Enable foreign key constraints in SQLite which are crucial to prevent programmer errors on
	// our side.
	_, err := database.Exec(
		`PRAGMA journal_mode=WAL;
		PRAGMA temp_store=1;
		PRAGMA foreign_keys=ON;`,
	)
	if err != nil {
		return err
	}

	_, err = database.Exec(
		`CREATE TABLE IF NOT EXISTS torrents (
			id             INTEGER PRIMARY KEY,
			info_hash      BLOB NOT NULL UNIQUE,
			name           TEXT NOT NULL,
			total_size     INTEGER NOT NULL CHECK(total_size > 0),
			discovered_on  INTEGER NOT NULL CHECK(discovered_on > 0)
		);

		CREATE INDEX IF NOT EXISTS info_hash_index ON torrents (info_hash);

		CREATE TABLE IF NOT EXISTS files (
			id          INTEGER PRIMARY KEY,
			torrent_id  INTEGER REFERENCES torrents ON DELETE CASCADE ON UPDATE RESTRICT,
			size        INTEGER NOT NULL,
			path        TEXT NOT NULL
		);`,
	)
	if err != nil {
		return err
	}

	return nil
}

func setupMySQLDatabase(database *sql.DB) error {
	_, err := database.Exec(
		`CREATE TABLE IF NOT EXISTS torrents ("
			id		INTEGER PRIMARY KEY AUTO_INCREMENT,
			info_hash	BINARY(20) NOT NULL UNIQUE,
			name		VARCHAR(255) NOT NULL,
			total_size	BIGINT UNSIGNED NOT NULL,
			discovered_on	INTEGER UNSIGNED NOT NULL
		);

		ALTER TABLE torrents ADD INDEX info_hash_index (info_hash);

		CREATE TABLE IF NOT EXISTS files (
			id		INTEGER PRIMARY KEY AUTO_INCREMENT,
			torrent_id	INTEGER REFERENCES torrents (id) ON DELETE CASCADE ON UPDATE RESTRICT,
			size		BIGINT NOT NULL,
			path		TEXT NOT NULL
		);`,
	)

	if err != nil {
		return err
	}

	return nil
}
<commit_msg>magneticod: use mysql in strict mode, increase name column to 1024 characters<commit_after>package main

import (
	"fmt"
	"database/sql"
	"net/url"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/mattn/go-sqlite3"
	"go.uber.org/zap"

	"magneticod/bittorrent"

	"path"
	"os"
)

type engineType uint8

const (
	SQLITE     engineType = 0
	POSTGRESQL            = 1
	MYSQL                 = 2
)

type Database struct {
	database    *sql.DB
	engine      engineType
	newTorrents chan bittorrent.Metadata
}

// NewDatabase creates a new Database.
//
// url either starts with "sqlite:" or "postgresql:"
func NewDatabase(rawurl string) (*Database, error) {
	db := Database{}

	dbURL, err := url.Parse(rawurl)
	if err != nil {
		return nil, err
	}

	switch dbURL.Scheme {
	case "sqlite":
		db.engine = SQLITE
		// All this pain is to make sure that an empty file exist (in case the database is not there
		// yet) so that sql.Open won't fail.
		dbDir, _ := path.Split(dbURL.Path)
		if err := os.MkdirAll(dbDir, 0755); err != nil {
			return nil, fmt.Errorf("for directory `%s`:  %s", dbDir, err.Error())
		}
		f, err := os.OpenFile(dbURL.Path, os.O_CREATE, 0666)
		if err != nil {
			return nil, fmt.Errorf("for file `%s`: %s", dbURL.Path, err.Error())
		}
		if err := f.Sync(); err != nil {
			return nil, fmt.Errorf("for file `%s`: %s", dbURL.Path, err.Error())
		}
		if err := f.Close(); err != nil {
			return nil, fmt.Errorf("for file `%s`: %s", dbURL.Path, err.Error())
		}
		db.database, err = sql.Open("sqlite3", dbURL.RawPath)

	case "postgresql":
		db.engine = POSTGRESQL
		db.database, err = sql.Open("postgresql", rawurl)

	case "mysql":
		db.engine = MYSQL
		db.database, err = sql.Open("mysql", rawurl)

	default:
		return nil, fmt.Errorf("unknown URI scheme (or malformed URI)!")
	}

	// Check for errors from sql.Open()
	if err != nil {
		return nil, fmt.Errorf("sql.Open()!  %s", err.Error())
	}

	if err = db.database.Ping(); err != nil {
		return nil, fmt.Errorf("DB.Ping()!  %s", err.Error())
	}

	if err := db.setupDatabase(); err != nil {
		return nil, fmt.Errorf("setupDatabase()!  %s", err.Error())
	}

	db.newTorrents = make(chan bittorrent.Metadata, 10)

	return &db, nil
}

// AddNewTorrent adds a new torrent to the *queue* to be flushed to the persistent database.
func (db *Database) AddNewTorrent(torrent bittorrent.Metadata) error {
	for {
		select {
		case db.newTorrents <- torrent:
			return nil
		default:
			// newTorrents queue was full: flush and try again and again (and again)...
			err := db.flushNewTorrents()
			if err != nil {
				return err
			}
			continue
		}
	}
}

func (db *Database) flushNewTorrents() error {
	tx, err := db.database.Begin()
	if err != nil {
		return fmt.Errorf("sql.DB.Begin()!  %s", err.Error())
	}

	var nTorrents, nFiles uint
	for torrent := range db.newTorrents {
		res, err := tx.Exec("INSERT INTO torrents (info_hash, name, total_size, discovered_on) VALUES ($1, $2, $3, $4);",
			torrent.InfoHash, torrent.Name, torrent.TotalSize, torrent.DiscoveredOn)
		if err != nil {
			ourError := fmt.Errorf("error while INSERTing INTO torrents!  %s", err.Error())
			if err := tx.Rollback(); err != nil {
				return fmt.Errorf("%s\tmeanwhile, could not rollback the current transaction either!  %s", ourError.Error(), err.Error())
			}
			return ourError
		}
		var lastInsertId int64
		if lastInsertId, err = res.LastInsertId(); err != nil {
			return fmt.Errorf("sql.Result.LastInsertId()!  %s", err.Error())
		}

		for _, file := range torrent.Files {
			_, err := tx.Exec("INSERT INTO files (torrent_id, size, path) VALUES($1, $2, $3);",
			lastInsertId, file.Length, file.Path)
			if err != nil {
				ourError := fmt.Errorf("error while INSERTing INTO files!  %s", err.Error())
				if err := tx.Rollback(); err != nil {
					return fmt.Errorf("%s\tmeanwhile, could not rollback the current transaction either!  %s", ourError.Error(), err.Error())
				}
				return ourError
			}
			nFiles++
		}
		nTorrents++
	}

	err = tx.Commit()
	if err != nil {
		return fmt.Errorf("sql.Tx.Commit()!  %s", err.Error())
	}

	zap.L().Sugar().Infof("%d torrents (%d files) are flushed to the database successfully.",
		nTorrents, nFiles)
	return nil
}

func (db *Database) Close() {
	// Be careful to not to get into an infinite loop. =)
	db.database.Close()
}

func (db *Database) setupDatabase() error {
	switch db.engine {
	case SQLITE:
		return setupSqliteDatabase(db.database)

	case POSTGRESQL:
		zap.L().Fatal("setupDatabase() is not implemented for PostgreSQL yet!")

	case MYSQL:
		return setupMySQLDatabase(db.database)

	default:
		zap.L().Sugar().Fatalf("Unknown database engine value %d! (programmer error)", db.engine)
	}

	return nil
}

func setupSqliteDatabase(database *sql.DB) error {
	// Enable Write-Ahead Logging for SQLite as "WAL provides more concurrency as readers do not
	// block writers and a writer does not block readers. Reading and writing can proceed
	// concurrently."
	// Caveats:
	//   * Might be unsupported by OSes other than Windows and UNIXes.
	//   * Does not work over a network filesystem.
	//   * Transactions that involve changes against multiple ATTACHed databases are not atomic
	//     across all databases as a set.
	// See: https://www.sqlite.org/wal.html
	//
	// Force SQLite to use disk, instead of memory, for all temporary files to reduce the memory
	// footprint.
	//
	// Enable foreign key constraints in SQLite which are crucial to prevent programmer errors on
	// our side.
	_, err := database.Exec(
		`PRAGMA journal_mode=WAL;
		PRAGMA temp_store=1;
		PRAGMA foreign_keys=ON;`,
	)
	if err != nil {
		return err
	}

	_, err = database.Exec(
		`CREATE TABLE IF NOT EXISTS torrents (
			id             INTEGER PRIMARY KEY,
			info_hash      BLOB NOT NULL UNIQUE,
			name           TEXT NOT NULL,
			total_size     INTEGER NOT NULL CHECK(total_size > 0),
			discovered_on  INTEGER NOT NULL CHECK(discovered_on > 0)
		);

		CREATE INDEX IF NOT EXISTS info_hash_index ON torrents (info_hash);

		CREATE TABLE IF NOT EXISTS files (
			id          INTEGER PRIMARY KEY,
			torrent_id  INTEGER REFERENCES torrents ON DELETE CASCADE ON UPDATE RESTRICT,
			size        INTEGER NOT NULL,
			path        TEXT NOT NULL
		);`,
	)
	if err != nil {
		return err
	}

	return nil
}

func setupMySQLDatabase(database *sql.DB) error {
	// Set strict mode to prevent silent truncation
	_, err := database.Exec(`SET SESSION SQL_MODE = 'STRICT_ALL_TABLES';`)
	if err != nil {
		return err
	}

	_, err = database.Exec(
		`CREATE TABLE IF NOT EXISTS torrents ("
			id		INTEGER PRIMARY KEY AUTO_INCREMENT,
			info_hash	BINARY(20) NOT NULL UNIQUE,
			name		VARCHAR(1024) NOT NULL,
			total_size	BIGINT UNSIGNED NOT NULL,
			discovered_on	INTEGER UNSIGNED NOT NULL
		);

		ALTER TABLE torrents ADD INDEX info_hash_index (info_hash);

		CREATE TABLE IF NOT EXISTS files (
			id		INTEGER PRIMARY KEY AUTO_INCREMENT,
			torrent_id	INTEGER REFERENCES torrents (id) ON DELETE CASCADE ON UPDATE RESTRICT,
			size		BIGINT NOT NULL,
			path		TEXT NOT NULL
		);`,
	)

	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>