<commit_before>/*

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 kubectl

import (
	"fmt"
	"strings"
	"time"

	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/client-go/kubernetes"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"sigs.k8s.io/multi-tenancy/benchmarks/kubectl-mtb/internal/reporter"
	"sigs.k8s.io/multi-tenancy/benchmarks/kubectl-mtb/pkg/benchmark"
	"sigs.k8s.io/multi-tenancy/benchmarks/kubectl-mtb/test"
	"sigs.k8s.io/multi-tenancy/benchmarks/kubectl-mtb/test/utils"
)

var (
	tenant          string
	tenantNamespace string
	k8sClient       *kubernetes.Clientset
	tenantClient    *kubernetes.Clientset
)

var runCmd = &cobra.Command{
	Use:   "run <resource>",
	Short: "Run the Multi-Tenancy Benchmarks",

	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 0 {
			return fmt.Errorf("Please specify a resource")
		}
		if !supportedResourceNames.Has(args[0]) {
			return fmt.Errorf("Please specify a valid resource")
		}
		err := validateFlags(cmd)
		if err != nil {
			return err
		}

		filterBenchmarks(cmd)

		return nil
	},

	Run: func(cmd *cobra.Command, args []string) {
		cmdutil.CheckErr(runTests(cmd, args))
	},
}

func initConfig() error {
	kubecfgFlags := genericclioptions.NewConfigFlags(false)

	config, err := kubecfgFlags.ToRESTConfig()
	if err != nil {
		return err
	}

	// create the K8s clientset
	k8sClient, err = kubernetes.NewForConfig(config)
	if err != nil {
		return err
	}

	tenantConfig := config
	tenantConfig.Impersonate.UserName = tenant

	// create the tenant clientset
	tenantClient, err = kubernetes.NewForConfig(tenantConfig)
	if err != nil {
		return err
	}

	return err
}

func reportSuiteWillBegin(suiteSummary *reporter.SuiteSummary, reportersArray []reporter.Reporter) {
	for _, reporter := range reportersArray {
		reporter.SuiteWillBegin(suiteSummary)
	}
}

func reportTestWillRun(testSummary *reporter.TestSummary, reportersArray []reporter.Reporter) {
	for _, reporter := range reportersArray {
		reporter.TestWillRun(testSummary)
	}
}

func reportSuiteDidEnd(suiteSummary *reporter.SuiteSummary, reportersArray []reporter.Reporter) {
	for _, reporter := range reportersArray {
		reporter.SuiteDidEnd(suiteSummary)
	}
}

func removeBenchmarksWithIDs(ids []string) {
	temp := []*benchmark.Benchmark{}
		for _, benchmark := range benchmarks {
			found := false
			for _, id := range ids {
				if(benchmark.ID == id) {
					found = true
				}
			}

			if !found {
				temp = append(temp, benchmark)
			}
		}
	benchmarks = temp
}

// Validation of the flag inputs
func validateFlags(cmd *cobra.Command) error {
	tenant, _ = cmd.Flags().GetString("as")
	if tenant == "" {
		return fmt.Errorf("username must be set via --as")
	}

	tenantNamespace, _ = cmd.Flags().GetString("namespace")
	if tenantNamespace == "" {
		return fmt.Errorf("tenant namespace must be set via --namespace or -n")
	}

	err := initConfig()
	if err != nil {
		return err
	}

	resource := utils.GroupResource{
		APIGroup: "",
		APIResource: metav1.APIResource{
			Name: "namespaces",
		},
		ResourceName: tenantNamespace,
	}
	// checks if tenant-admin and tenant namespace are valid
	access, _, err := utils.RunAccessCheck(tenantClient, "", resource, "get")
	if err != nil {
		return err
	}
	if !access {
		return fmt.Errorf("Make sure you have entered valid tenant-admin and tenant namespace. ")
	}
	return nil
}

func runTests(cmd *cobra.Command, args []string) error {

	// Get reporters from the user
	reporterFlag, _ := cmd.Flags().GetString("out")
	reporters := strings.Split(reporterFlag, ",")
	reportersArray, err := reporter.GetReporters(reporters)
	if err != nil {
		return err
	}

	// Get benchmark ids from the user to skip them
	skipFlag, _ := cmd.Flags().GetString("skip")
	skipIDs := strings.Split(skipFlag, ",")
	removeBenchmarksWithIDs(skipIDs)

	suiteSummary := &reporter.SuiteSummary{
		Suite:                test.BenchmarkSuite,
		NumberOfTotalTests:   len(benchmarks),
		TenantAdminNamespace: tenantNamespace,
	}

	suiteStartTime := time.Now()
	reportSuiteWillBegin(suiteSummary, reportersArray)

	for _, b := range benchmarks {

		ts := &reporter.TestSummary{
			Benchmark: b,
		}

		err := ts.SetDefaults()
		if err != nil {
			return err
		}

		startTest := time.Now()

		//Run Prerun
		err = b.PreRun(tenantNamespace, k8sClient, tenantClient)
		if err != nil {
			suiteSummary.NumberOfFailedValidations++
			ts.Validation = false
			ts.ValidationError = err
		}

		// Check PreRun status
		if ts.Validation {
			err = b.Run(tenantNamespace, k8sClient, tenantClient)
			if err != nil {
				suiteSummary.NumberOfFailedTests++
				ts.Test = false
				ts.TestError = err
			} else {
				suiteSummary.NumberOfPassedTests++
			}
		}
		elapsed := time.Since(startTest)
		ts.RunTime = elapsed
		reportTestWillRun(ts, reportersArray)
	}

	suiteElapsedTime := time.Since(suiteStartTime)
	suiteSummary.RunTime = suiteElapsedTime
	suiteSummary.NumberOfSkippedTests = test.BenchmarkSuite.Totals() - len(benchmarks)
	reportSuiteDidEnd(suiteSummary, reportersArray)

	return nil
}

func newRunCmd() *cobra.Command {
	runCmd.Flags().StringP("namespace", "n", "", "(required) tenant namespace")
	runCmd.Flags().String("as", "", "(required) user name to impersonate")
	runCmd.Flags().StringP("out", "o", "default", "(optional) output reporters (default, policyreport)")
	runCmd.Flags().StringP("skip", "s", "", "(optional) benchmark IDs to skip")

	return runCmd
}
<commit_msg>changed validation check<commit_after>/*

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 kubectl

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/client-go/kubernetes"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"sigs.k8s.io/multi-tenancy/benchmarks/kubectl-mtb/internal/reporter"
	"sigs.k8s.io/multi-tenancy/benchmarks/kubectl-mtb/pkg/benchmark"
	"sigs.k8s.io/multi-tenancy/benchmarks/kubectl-mtb/test"
)

var (
	tenant          string
	tenantNamespace string
	k8sClient       *kubernetes.Clientset
	tenantClient    *kubernetes.Clientset
)

var runCmd = &cobra.Command{
	Use:   "run <resource>",
	Short: "Run the Multi-Tenancy Benchmarks",

	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 0 {
			return fmt.Errorf("Please specify a resource")
		}
		if !supportedResourceNames.Has(args[0]) {
			return fmt.Errorf("Please specify a valid resource")
		}
		err := validateFlags(cmd)
		if err != nil {
			return err
		}

		filterBenchmarks(cmd)

		return nil
	},

	Run: func(cmd *cobra.Command, args []string) {
		cmdutil.CheckErr(runTests(cmd, args))
	},
}

func initConfig() error {
	kubecfgFlags := genericclioptions.NewConfigFlags(false)

	config, err := kubecfgFlags.ToRESTConfig()
	if err != nil {
		return err
	}

	// create the K8s clientset
	k8sClient, err = kubernetes.NewForConfig(config)
	if err != nil {
		return err
	}

	tenantConfig := config
	tenantConfig.Impersonate.UserName = tenant

	// create the tenant clientset
	tenantClient, err = kubernetes.NewForConfig(tenantConfig)
	if err != nil {
		return err
	}

	return err
}

func reportSuiteWillBegin(suiteSummary *reporter.SuiteSummary, reportersArray []reporter.Reporter) {
	for _, reporter := range reportersArray {
		reporter.SuiteWillBegin(suiteSummary)
	}
}

func reportTestWillRun(testSummary *reporter.TestSummary, reportersArray []reporter.Reporter) {
	for _, reporter := range reportersArray {
		reporter.TestWillRun(testSummary)
	}
}

func reportSuiteDidEnd(suiteSummary *reporter.SuiteSummary, reportersArray []reporter.Reporter) {
	for _, reporter := range reportersArray {
		reporter.SuiteDidEnd(suiteSummary)
	}
}

func removeBenchmarksWithIDs(ids []string) {
	temp := []*benchmark.Benchmark{}
		for _, benchmark := range benchmarks {
			found := false
			for _, id := range ids {
				if(benchmark.ID == id) {
					found = true
				}
			}

			if !found {
				temp = append(temp, benchmark)
			}
		}
	benchmarks = temp
}

// Validation of the flag inputs
func validateFlags(cmd *cobra.Command) error {
	tenant, _ = cmd.Flags().GetString("as")
	if tenant == "" {
		return fmt.Errorf("username must be set via --as")
	}

	tenantNamespace, _ = cmd.Flags().GetString("namespace")
	if tenantNamespace == "" {
		return fmt.Errorf("tenant namespace must be set via --namespace or -n")
	}

	err := initConfig()
	if err != nil {
		return err
	}

	_, err = k8sClient.CoreV1().Namespaces().Get(context.TODO(), tenantNamespace, metav1.GetOptions{})
	if err != nil {
		return fmt.Errorf("tenantnamespace is not a valid namespace")
	}

	return nil
}

func runTests(cmd *cobra.Command, args []string) error {

	// Get reporters from the user
	reporterFlag, _ := cmd.Flags().GetString("out")
	reporters := strings.Split(reporterFlag, ",")
	reportersArray, err := reporter.GetReporters(reporters)
	if err != nil {
		return err
	}

	// Get benchmark ids from the user to skip them
	skipFlag, _ := cmd.Flags().GetString("skip")
	skipIDs := strings.Split(skipFlag, ",")
	removeBenchmarksWithIDs(skipIDs)

	suiteSummary := &reporter.SuiteSummary{
		Suite:                test.BenchmarkSuite,
		NumberOfTotalTests:   len(benchmarks),
		TenantAdminNamespace: tenantNamespace,
	}

	suiteStartTime := time.Now()
	reportSuiteWillBegin(suiteSummary, reportersArray)

	for _, b := range benchmarks {

		ts := &reporter.TestSummary{
			Benchmark: b,
		}

		err := ts.SetDefaults()
		if err != nil {
			return err
		}

		startTest := time.Now()

		//Run Prerun
		err = b.PreRun(tenantNamespace, k8sClient, tenantClient)
		if err != nil {
			suiteSummary.NumberOfFailedValidations++
			ts.Validation = false
			ts.ValidationError = err
		}

		// Check PreRun status
		if ts.Validation {
			err = b.Run(tenantNamespace, k8sClient, tenantClient)
			if err != nil {
				suiteSummary.NumberOfFailedTests++
				ts.Test = false
				ts.TestError = err
			} else {
				suiteSummary.NumberOfPassedTests++
			}
		}
		elapsed := time.Since(startTest)
		ts.RunTime = elapsed
		reportTestWillRun(ts, reportersArray)
	}

	suiteElapsedTime := time.Since(suiteStartTime)
	suiteSummary.RunTime = suiteElapsedTime
	suiteSummary.NumberOfSkippedTests = test.BenchmarkSuite.Totals() - len(benchmarks)
	reportSuiteDidEnd(suiteSummary, reportersArray)

	return nil
}

func newRunCmd() *cobra.Command {
	runCmd.Flags().StringP("namespace", "n", "", "(required) tenant namespace")
	runCmd.Flags().String("as", "", "(required) user name to impersonate")
	runCmd.Flags().StringP("out", "o", "default", "(optional) output reporters (default, policyreport)")
	runCmd.Flags().StringP("skip", "s", "", "(optional) benchmark IDs to skip")

	return runCmd
}
<|endoftext|>


<commit_before>package util

import (
	"testing"
	"time"

	"github.com/stretchr/testify/suite"
)

type Suite struct {
	suite.Suite
}

func (suite *Suite) TestNewTimePeriod() {
	timezone, err := time.LoadLocation("Australia/Brisbane")
	suite.Require().NoError(err)

	for _, tt := range []struct {
		from     time.Time
		to       time.Time
		expected TimePeriod
	}{
		// when it's already normalized
		{
			time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
			time.Date(0, 0, 0, 16, 04, 05, 06, time.UTC),
			TimePeriod{
				From: time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
				To:   time.Date(0, 0, 0, 16, 04, 05, 06, time.UTC),
			},
		},
		// it normalizes to very first day
		{
			time.Date(1869, 9, 24, 15, 04, 05, 06, time.UTC),
			time.Date(1869, 9, 24, 16, 04, 05, 06, time.UTC),
			TimePeriod{
				From: time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
				To:   time.Date(0, 0, 0, 16, 04, 05, 06, time.UTC),
			},
		},
		// it ignores the timezone
		{
			time.Date(1869, 9, 24, 15, 04, 05, 06, timezone),
			time.Date(1869, 9, 24, 16, 04, 05, 06, timezone),
			TimePeriod{
				From: time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
				To:   time.Date(0, 0, 0, 16, 04, 05, 06, time.UTC),
			},
		},
	} {
		suite.Equal(tt.expected, NewTimePeriod(tt.from, tt.to))
	}
}

func (suite *Suite) TestTimePeriodIncludes() {
	atTheMoment := NewTimePeriod(
		time.Now().Add(-1*time.Minute),
		time.Now().Add(+1*time.Minute),
	)
	midnight := NewTimePeriod(
		time.Date(1869, 9, 23, 23, 00, 00, 00, time.UTC),
		time.Date(1869, 9, 24, 01, 00, 00, 00, time.UTC),
	)
	now := time.Now()

	for _, tt := range []struct {
		pointInTime time.Time
		timeOfDay   TimePeriod
		expected    bool
	}{
		// it's included
		{
			now,
			atTheMoment,
			true,
		},
		// it's one day before
		{
			now.Add(-24 * time.Hour),
			atTheMoment,
			true,
		},
		// it's one day after
		{
			now.Add(+24 * time.Hour),
			atTheMoment,
			true,
		},
		// it's just before
		{
			now.Add(-2 * time.Minute),
			atTheMoment,
			false,
		},
		// it's just after
		{
			now.Add(+2 * time.Minute),
			atTheMoment,
			false,
		},
		// it's slightly inside before day switch
		{
			time.Date(1869, 9, 23, 23, 30, 00, 00, time.UTC),
			midnight,
			true,
		},
		// it's slightly inside after day switch
		{
			time.Date(1869, 9, 24, 00, 30, 00, 00, time.UTC),
			midnight,
			true,
		},
		// it's just before
		{
			time.Date(1869, 9, 23, 22, 30, 00, 00, time.UTC),
			midnight,
			false,
		},
		// it's just after
		{
			time.Date(1869, 9, 24, 01, 30, 00, 00, time.UTC),
			midnight,
			false,
		},
		// it's exactly matching a point in time
		{
			now,
			TimePeriod{From: TimeOfDay(now), To: TimeOfDay(now)},
			true,
		},
		// it's right after exactly matching a point in time
		{
			now.Add(+1 * time.Second),
			TimePeriod{From: TimeOfDay(now), To: TimeOfDay(now)},
			false,
		},
		// it's right before exactly matching a point in time
		{
			now.Add(-1 * time.Second),
			TimePeriod{From: TimeOfDay(now), To: TimeOfDay(now)},
			false,
		},
	} {
		suite.Equal(tt.expected, tt.timeOfDay.Includes(tt.pointInTime))
	}
}

func (suite *Suite) TestTimePeriodString() {
	for _, tt := range []struct {
		given    TimePeriod
		expected string
	}{
		// empty time period
		{
			TimePeriod{},
			"00:00-00:00",
		},
		// simple time period
		{
			TimePeriod{
				From: time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
				To:   time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
			},
			"08:00-16:00",
		},
		// time period across days
		{
			TimePeriod{
				From: time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
				To:   time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
			},
			"16:00-08:00",
		},
	} {
		suite.Equal(tt.expected, tt.given.String())
	}
}

func (suite *Suite) TestTimeOfDay() {
	timezone, err := time.LoadLocation("Australia/Brisbane")
	suite.Require().NoError(err)

	for _, tt := range []struct {
		pointInTime time.Time
		expected    time.Time
	}{
		// strips of any day information
		{
			time.Date(1869, 9, 24, 15, 04, 05, 06, time.UTC),
			time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
		},
		// it normalizes to UTC timezone
		{
			time.Date(1869, 9, 24, 15, 04, 05, 06, timezone),
			time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
		},
	} {
		suite.Equal(tt.expected, TimeOfDay(tt.pointInTime))
	}
}

func (suite *Suite) TestParseWeekdays() {
	for _, tt := range []struct {
		given    string
		expected []time.Weekday
	}{
		// empty string
		{
			"",
			[]time.Weekday{},
		},
		// single weekday
		{
			"sat",
			[]time.Weekday{time.Saturday},
		},
		// multiple weekdays
		{
			"sat,sun",
			[]time.Weekday{time.Saturday, time.Sunday},
		},
		// case-insensitive
		{
			"SaT,SUn",
			[]time.Weekday{time.Saturday, time.Sunday},
		},
		// ignore whitespace
		{
			" sat , sun ",
			[]time.Weekday{time.Saturday, time.Sunday},
		},
		// ignore unknown weekdays
		{
			"sat,unknown,sun",
			[]time.Weekday{time.Saturday, time.Sunday},
		},
		// deal with all kinds at the same time
		{
			"Fri, sat ,,,,  ,foobar,tue",
			[]time.Weekday{time.Friday, time.Saturday, time.Tuesday},
		},
	} {
		suite.Equal(tt.expected, ParseWeekdays(tt.given))
	}
}

func (suite *Suite) TestParseTimePeriods() {
	for _, tt := range []struct {
		given    string
		expected []TimePeriod
	}{
		// empty time period string
		{
			"",
			[]TimePeriod{},
		},
		// single range string
		{
			"08:00-16:00",
			[]TimePeriod{
				{
					From: time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
				},
			},
		},
		// multiple ranges string
		{
			"08:00-16:00,20:00-22:00",
			[]TimePeriod{
				{
					From: time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
				},
				{
					From: time.Date(0, 0, 0, 20, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 22, 0, 0, 0, time.UTC),
				},
			},
		},
		// string containing whitespace
		{
			" 08:00 - 16:00 ,, , 20:00 - 22:00 ",
			[]TimePeriod{
				{
					From: time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
				},
				{
					From: time.Date(0, 0, 0, 20, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 22, 0, 0, 0, time.UTC),
				},
			},
		},
	} {
		timePeriods, err := ParseTimePeriods(tt.given)
		suite.Require().NoError(err)

		suite.Equal(tt.expected, timePeriods)
	}
}

func (suite *Suite) TestParseDates() {
	for _, tt := range []struct {
		given    string
		expected []time.Time
	}{
		// empty string
		{
			"",
			[]time.Time{},
		},
		// single date
		{
			"Apr 1",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
			},
		},
		// single date leaving out the space
		{
			"Apr1",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
			},
		},
		// multiple dates
		{
			"Apr 1,Dec 24",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
				time.Date(0, 12, 24, 0, 0, 0, 0, time.UTC),
			},
		},
		// case-insensitive
		{
			"apr 1,dEc 24",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
				time.Date(0, 12, 24, 0, 0, 0, 0, time.UTC),
			},
		},
		// ignore whitespace
		{
			" apr 1 , dec 24 ",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
				time.Date(0, 12, 24, 0, 0, 0, 0, time.UTC),
			},
		},
		// deal with all kinds at the same time
		{
			",Apr 1, dEc 24 ,,,,  ,jun08,,",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
				time.Date(0, 12, 24, 0, 0, 0, 0, time.UTC),
				time.Date(0, 6, 8, 0, 0, 0, 0, time.UTC),
			},
		},
	} {
		days, err := ParseDays(tt.given)
		suite.Require().NoError(err)

		suite.Equal(tt.expected, days)
	}
}

func (suite *Suite) TestFormatDays() {
	for _, tt := range []struct {
		given    []time.Time
		expected []string
	}{
		{
			[]time.Time{
				time.Date(1869, 9, 24, 15, 04, 05, 06, time.UTC),
			},
			[]string{"Sep24"},
		}, {
			[]time.Time{
				time.Date(1869, 9, 24, 15, 04, 05, 06, time.UTC),
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
			},
			[]string{"Sep24", "Apr 1"},
		},
	} {
		suite.Equal(tt.expected, FormatDays(tt.given))
	}
}

func TestSuite(t *testing.T) {
	suite.Run(t, new(Suite))
}
<commit_msg>chore: add test for NewNamespace helper function<commit_after>package util

import (
	"testing"
	"time"

	"github.com/stretchr/testify/suite"
)

type Suite struct {
	suite.Suite
}

func (suite *Suite) TestNewTimePeriod() {
	timezone, err := time.LoadLocation("Australia/Brisbane")
	suite.Require().NoError(err)

	for _, tt := range []struct {
		from     time.Time
		to       time.Time
		expected TimePeriod
	}{
		// when it's already normalized
		{
			time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
			time.Date(0, 0, 0, 16, 04, 05, 06, time.UTC),
			TimePeriod{
				From: time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
				To:   time.Date(0, 0, 0, 16, 04, 05, 06, time.UTC),
			},
		},
		// it normalizes to very first day
		{
			time.Date(1869, 9, 24, 15, 04, 05, 06, time.UTC),
			time.Date(1869, 9, 24, 16, 04, 05, 06, time.UTC),
			TimePeriod{
				From: time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
				To:   time.Date(0, 0, 0, 16, 04, 05, 06, time.UTC),
			},
		},
		// it ignores the timezone
		{
			time.Date(1869, 9, 24, 15, 04, 05, 06, timezone),
			time.Date(1869, 9, 24, 16, 04, 05, 06, timezone),
			TimePeriod{
				From: time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
				To:   time.Date(0, 0, 0, 16, 04, 05, 06, time.UTC),
			},
		},
	} {
		suite.Equal(tt.expected, NewTimePeriod(tt.from, tt.to))
	}
}

func (suite *Suite) TestTimePeriodIncludes() {
	atTheMoment := NewTimePeriod(
		time.Now().Add(-1*time.Minute),
		time.Now().Add(+1*time.Minute),
	)
	midnight := NewTimePeriod(
		time.Date(1869, 9, 23, 23, 00, 00, 00, time.UTC),
		time.Date(1869, 9, 24, 01, 00, 00, 00, time.UTC),
	)
	now := time.Now()

	for _, tt := range []struct {
		pointInTime time.Time
		timeOfDay   TimePeriod
		expected    bool
	}{
		// it's included
		{
			now,
			atTheMoment,
			true,
		},
		// it's one day before
		{
			now.Add(-24 * time.Hour),
			atTheMoment,
			true,
		},
		// it's one day after
		{
			now.Add(+24 * time.Hour),
			atTheMoment,
			true,
		},
		// it's just before
		{
			now.Add(-2 * time.Minute),
			atTheMoment,
			false,
		},
		// it's just after
		{
			now.Add(+2 * time.Minute),
			atTheMoment,
			false,
		},
		// it's slightly inside before day switch
		{
			time.Date(1869, 9, 23, 23, 30, 00, 00, time.UTC),
			midnight,
			true,
		},
		// it's slightly inside after day switch
		{
			time.Date(1869, 9, 24, 00, 30, 00, 00, time.UTC),
			midnight,
			true,
		},
		// it's just before
		{
			time.Date(1869, 9, 23, 22, 30, 00, 00, time.UTC),
			midnight,
			false,
		},
		// it's just after
		{
			time.Date(1869, 9, 24, 01, 30, 00, 00, time.UTC),
			midnight,
			false,
		},
		// it's exactly matching a point in time
		{
			now,
			TimePeriod{From: TimeOfDay(now), To: TimeOfDay(now)},
			true,
		},
		// it's right after exactly matching a point in time
		{
			now.Add(+1 * time.Second),
			TimePeriod{From: TimeOfDay(now), To: TimeOfDay(now)},
			false,
		},
		// it's right before exactly matching a point in time
		{
			now.Add(-1 * time.Second),
			TimePeriod{From: TimeOfDay(now), To: TimeOfDay(now)},
			false,
		},
	} {
		suite.Equal(tt.expected, tt.timeOfDay.Includes(tt.pointInTime))
	}
}

func (suite *Suite) TestTimePeriodString() {
	for _, tt := range []struct {
		given    TimePeriod
		expected string
	}{
		// empty time period
		{
			TimePeriod{},
			"00:00-00:00",
		},
		// simple time period
		{
			TimePeriod{
				From: time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
				To:   time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
			},
			"08:00-16:00",
		},
		// time period across days
		{
			TimePeriod{
				From: time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
				To:   time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
			},
			"16:00-08:00",
		},
	} {
		suite.Equal(tt.expected, tt.given.String())
	}
}

func (suite *Suite) TestTimeOfDay() {
	timezone, err := time.LoadLocation("Australia/Brisbane")
	suite.Require().NoError(err)

	for _, tt := range []struct {
		pointInTime time.Time
		expected    time.Time
	}{
		// strips of any day information
		{
			time.Date(1869, 9, 24, 15, 04, 05, 06, time.UTC),
			time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
		},
		// it normalizes to UTC timezone
		{
			time.Date(1869, 9, 24, 15, 04, 05, 06, timezone),
			time.Date(0, 0, 0, 15, 04, 05, 06, time.UTC),
		},
	} {
		suite.Equal(tt.expected, TimeOfDay(tt.pointInTime))
	}
}

func (suite *Suite) TestParseWeekdays() {
	for _, tt := range []struct {
		given    string
		expected []time.Weekday
	}{
		// empty string
		{
			"",
			[]time.Weekday{},
		},
		// single weekday
		{
			"sat",
			[]time.Weekday{time.Saturday},
		},
		// multiple weekdays
		{
			"sat,sun",
			[]time.Weekday{time.Saturday, time.Sunday},
		},
		// case-insensitive
		{
			"SaT,SUn",
			[]time.Weekday{time.Saturday, time.Sunday},
		},
		// ignore whitespace
		{
			" sat , sun ",
			[]time.Weekday{time.Saturday, time.Sunday},
		},
		// ignore unknown weekdays
		{
			"sat,unknown,sun",
			[]time.Weekday{time.Saturday, time.Sunday},
		},
		// deal with all kinds at the same time
		{
			"Fri, sat ,,,,  ,foobar,tue",
			[]time.Weekday{time.Friday, time.Saturday, time.Tuesday},
		},
	} {
		suite.Equal(tt.expected, ParseWeekdays(tt.given))
	}
}

func (suite *Suite) TestParseTimePeriods() {
	for _, tt := range []struct {
		given    string
		expected []TimePeriod
	}{
		// empty time period string
		{
			"",
			[]TimePeriod{},
		},
		// single range string
		{
			"08:00-16:00",
			[]TimePeriod{
				{
					From: time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
				},
			},
		},
		// multiple ranges string
		{
			"08:00-16:00,20:00-22:00",
			[]TimePeriod{
				{
					From: time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
				},
				{
					From: time.Date(0, 0, 0, 20, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 22, 0, 0, 0, time.UTC),
				},
			},
		},
		// string containing whitespace
		{
			" 08:00 - 16:00 ,, , 20:00 - 22:00 ",
			[]TimePeriod{
				{
					From: time.Date(0, 0, 0, 8, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 16, 0, 0, 0, time.UTC),
				},
				{
					From: time.Date(0, 0, 0, 20, 0, 0, 0, time.UTC),
					To:   time.Date(0, 0, 0, 22, 0, 0, 0, time.UTC),
				},
			},
		},
	} {
		timePeriods, err := ParseTimePeriods(tt.given)
		suite.Require().NoError(err)

		suite.Equal(tt.expected, timePeriods)
	}
}

func (suite *Suite) TestParseDates() {
	for _, tt := range []struct {
		given    string
		expected []time.Time
	}{
		// empty string
		{
			"",
			[]time.Time{},
		},
		// single date
		{
			"Apr 1",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
			},
		},
		// single date leaving out the space
		{
			"Apr1",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
			},
		},
		// multiple dates
		{
			"Apr 1,Dec 24",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
				time.Date(0, 12, 24, 0, 0, 0, 0, time.UTC),
			},
		},
		// case-insensitive
		{
			"apr 1,dEc 24",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
				time.Date(0, 12, 24, 0, 0, 0, 0, time.UTC),
			},
		},
		// ignore whitespace
		{
			" apr 1 , dec 24 ",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
				time.Date(0, 12, 24, 0, 0, 0, 0, time.UTC),
			},
		},
		// deal with all kinds at the same time
		{
			",Apr 1, dEc 24 ,,,,  ,jun08,,",
			[]time.Time{
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
				time.Date(0, 12, 24, 0, 0, 0, 0, time.UTC),
				time.Date(0, 6, 8, 0, 0, 0, 0, time.UTC),
			},
		},
	} {
		days, err := ParseDays(tt.given)
		suite.Require().NoError(err)

		suite.Equal(tt.expected, days)
	}
}

func (suite *Suite) TestFormatDays() {
	for _, tt := range []struct {
		given    []time.Time
		expected []string
	}{
		{
			[]time.Time{
				time.Date(1869, 9, 24, 15, 04, 05, 06, time.UTC),
			},
			[]string{"Sep24"},
		}, {
			[]time.Time{
				time.Date(1869, 9, 24, 15, 04, 05, 06, time.UTC),
				time.Date(0, 4, 1, 0, 0, 0, 0, time.UTC),
			},
			[]string{"Sep24", "Apr 1"},
		},
	} {
		suite.Equal(tt.expected, FormatDays(tt.given))
	}
}

func (suite *Suite) TestNewNamespace() {
	namespace := NewNamespace("name")

	suite.Equal("name", namespace.Name)
	suite.Equal("name", namespace.Labels["env"])
}

func TestSuite(t *testing.T) {
	suite.Run(t, new(Suite))
}
<|endoftext|>


<commit_before>// Copyright 2016 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.

package ld

import (
	"cmd/internal/obj"
	"strings"
)

// mergestrings merges all go.string.* character data into a single symbol.
//
// Combining string data symbols reduces the total binary size and
// makes deduplication possible.
func mergestrings() {
	if Buildmode == BuildmodeShared {
		return
	}

	strs := make([]*LSym, 0, 256)
	seenStr := make(map[string]bool, 256)         // symbol name -> in strs slice
	relocsToStrs := make(map[*LSym][]*Reloc, 256) // string -> relocation to string
	size := 0                                     // number of bytes in all strings

	// Collect strings and relocations that point to strings.
	for _, s := range Ctxt.Allsym {
		if !s.Attr.Reachable() || s.Attr.Special() {
			continue
		}
		for i := range s.R {
			r := &s.R[i]
			if r.Sym == nil {
				continue
			}
			if !seenStr[r.Sym.Name] {
				if !strings.HasPrefix(r.Sym.Name, "go.string.") {
					continue
				}
				if strings.HasPrefix(r.Sym.Name, "go.string.hdr") {
					continue
				}
				strs = append(strs, r.Sym)
				seenStr[r.Sym.Name] = true
				size += len(r.Sym.P)
			}
			relocsToStrs[r.Sym] = append(relocsToStrs[r.Sym], r)
		}
	}

	// Put all string data into a single symbol and update the relocations.
	alldata := Linklookup(Ctxt, "go.string.alldata", 0)
	alldata.Type = obj.SGOSTRING
	alldata.Attr |= AttrReachable
	alldata.P = make([]byte, 0, size)
	for _, str := range strs {
		off := len(alldata.P)
		alldata.P = append(alldata.P, str.P...)
		// Architectures with Minalign > 1 cannot have relocations pointing
		// to arbitrary locations, so make sure each string is appropriately
		// aligned.
		for r := len(alldata.P) % Thearch.Minalign; r > 0; r-- {
			alldata.P = append(alldata.P, 0)
		}
		str.Attr.Set(AttrReachable, false)
		for _, r := range relocsToStrs[str] {
			r.Add += int64(off)
			r.Sym = alldata
		}
	}
	alldata.Size = int64(len(alldata.P))
}
<commit_msg>cmd/link: deduplicate read-only string data<commit_after>// Copyright 2016 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.

package ld

import (
	"bytes"
	"cmd/internal/obj"
	"index/suffixarray"
	"sort"
	"strings"
)

// mergestrings merges all go.string.* character data into a single symbol.
//
// Combining string data symbols reduces the total binary size and
// makes deduplication possible.
func mergestrings() {
	if Buildmode == BuildmodeShared {
		return
	}

	strs := make([]*LSym, 0, 256)
	seenStr := make(map[string]bool, 256)         // symbol name -> in strs slice
	relocsToStrs := make(map[*LSym][]*Reloc, 256) // string -> relocation to string
	size := 0                                     // number of bytes in all strings

	// Collect strings and relocations that point to strings.
	for _, s := range Ctxt.Allsym {
		if !s.Attr.Reachable() || s.Attr.Special() {
			continue
		}
		for i := range s.R {
			r := &s.R[i]
			if r.Sym == nil {
				continue
			}
			if !seenStr[r.Sym.Name] {
				if !strings.HasPrefix(r.Sym.Name, "go.string.") {
					continue
				}
				if strings.HasPrefix(r.Sym.Name, "go.string.hdr") {
					continue
				}
				strs = append(strs, r.Sym)
				seenStr[r.Sym.Name] = true
				size += len(r.Sym.P)
			}
			relocsToStrs[r.Sym] = append(relocsToStrs[r.Sym], r)
		}
	}

	// Sort the strings, shortest first.
	//
	// Ordering by length lets us use the largest matching substring
	// index when there are multiple matches. This means we will not
	// use a substring of a string that we will later in the pass
	// mark as unreachable.
	sort.Sort(strSymsByLen(strs))

	// Build a suffix array.
	dataOff := make([]int, len(strs))
	data := make([]byte, 0, size)
	for i := range strs {
		dataOff[i] = len(data)
		data = append(data, strs[i].P...)
	}
	index := suffixarray.New(data)

	// Search for substring replacements.
	type replacement struct {
		str *LSym
		off int
	}
	replacements := make(map[*LSym]replacement)
	for i, s := range strs {
		results := index.Lookup(s.P, -1)
		if len(results) == 0 {
			continue
		}
		var res int
		for _, result := range results {
			if result > res {
				res = result
			}
		}
		var off int
		x := sort.SearchInts(dataOff, res)
		if x == len(dataOff) || dataOff[x] > res {
			x--
			off = res - dataOff[x]
		}
		if x == i {
			continue // found ourself
		}
		if len(s.P) > len(strs[x].P[off:]) {
			// Do not use substrings that match across strings.
			// In theory it is possible, but it would
			// complicate accounting for which future strings
			// are already used. It doesn't appear to be common
			// enough to do the extra work.
			continue
		}
		if off%Thearch.Minalign != 0 {
			continue // Cannot relcate to this substring.
		}
		replacements[s] = replacement{
			str: strs[x],
			off: off,
		}
	}

	// Put all string data into a single symbol and update the relocations.
	alldata := Linklookup(Ctxt, "go.string.alldata", 0)
	alldata.Type = obj.SGOSTRING
	alldata.Attr |= AttrReachable
	alldata.P = make([]byte, 0, size)
	for _, str := range strs {
		str.Attr.Set(AttrReachable, false)
		if rep, isReplaced := replacements[str]; isReplaced {
			// As strs is sorted, the replacement string
			// is always later in the strs range. Shift the
			// relocations to the replacement string symbol
			// and process then.
			relocs := relocsToStrs[rep.str]
			for _, r := range relocsToStrs[str] {
				r.Add += int64(rep.off)
				relocs = append(relocs, r)
			}
			relocsToStrs[rep.str] = relocs
			continue
		}

		off := len(alldata.P)
		alldata.P = append(alldata.P, str.P...)
		// Architectures with Minalign > 1 cannot have relocations pointing
		// to arbitrary locations, so make sure each string is appropriately
		// aligned.
		for r := len(alldata.P) % Thearch.Minalign; r > 0; r-- {
			alldata.P = append(alldata.P, 0)
		}
		for _, r := range relocsToStrs[str] {
			r.Add += int64(off)
			r.Sym = alldata
		}
	}
	alldata.Size = int64(len(alldata.P))
}

// strSymsByLen implements sort.Interface. It sorts *LSym by the length of P.
type strSymsByLen []*LSym

func (s strSymsByLen) Len() int      { return len(s) }
func (s strSymsByLen) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s strSymsByLen) Less(i, j int) bool {
	if len(s[i].P) == len(s[j].P) {
		return bytes.Compare(s[i].P, s[j].P) == -1
	}
	return len(s[i].P) < len(s[j].P)
}
<|endoftext|>


<commit_before>package knx

import (
	"container/list"

	"sync"

	"github.com/vapourismo/knx-go/knx/cemi"
	"github.com/vapourismo/knx-go/knx/proto"
)

// A RouterConfig determines certain properties of a Router.
type RouterConfig struct {
	// Specify how many sent messages to retain. This is important for when a router indicates that
	// it has lost some messages. If you do not expect to saturate the router, keep this low.
	RetainCount uint
}

// Default configuration elements
var (
	defaultRetainCount uint = 32

	DefaultRouterConfig = RouterConfig{
		defaultRetainCount,
	}
)

// checkRouterConfig validates the given RouterConfig.
func checkRouterConfig(config RouterConfig) RouterConfig {
	if config.RetainCount == 0 {
		config.RetainCount = defaultRetainCount
	}

	return config
}

// tryPushInbound sends the message through the channel. If the sending blocks, it will launch a
// goroutine which will do the sending.
func tryPushInbound(msg cemi.Message, inbound chan<- cemi.Message) {
	select {
	case inbound <- msg:

	default:
		go func() {
			// Since this goroutine decouples from the server goroutine, it might try to send when
			// the server closed the inbound channel. Sending to a closed channel will panic. But we
			// don't care, because cool guys don't look at explosions.
			defer func() { recover() }()
			inbound <- msg
		}()
	}
}

// A Router is a participant in a KNXnet/IP multicast group.
type Router struct {
	sock     Socket
	config   RouterConfig
	inbound  chan cemi.Message
	sendMu   sync.Mutex
	retainer *list.List
}

// serve listens for incoming routing-related packets.
func (router *Router) serve() {
	defer close(router.inbound)

	for msg := range router.sock.Inbound() {
		switch msg := msg.(type) {
		case *proto.RoutingInd:
			// Try to push it to the client without blocking this goroutine to long.
			tryPushInbound(msg.Payload, router.inbound)

		case *proto.RoutingBusy:
			// TODO: Inhibit sending for msg.WaitTime.

		case *proto.RoutingLost:
			// TODO: Resend the last msg.Count frames.
		}
	}
}

// NewRouter creates a new Router that joins the given multicast group.
func NewRouter(multicastAddress string, config RouterConfig) (*Router, error) {
	sock, err := NewRoutingSocket(multicastAddress)
	if err != nil {
		return nil, err
	}

	r := &Router{
		sock:     sock,
		config:   checkRouterConfig(config),
		inbound:  make(chan cemi.Message),
		retainer: list.New(),
	}

	go r.serve()

	return r, nil
}

// Send transmits a packet.
func (router *Router) Send(data cemi.Message) error {
	// We lock this before doing any sending so the server goroutine can adjust the flow control.
	router.sendMu.Lock()
	defer router.sendMu.Unlock()

	err := router.sock.Send(&proto.RoutingInd{Payload: data})

	if err == nil {
		// Store this for potential resending.
		router.retainer.PushBack(data)

		// We don't want to keep more messages than necessary. The overhead needs to be removed.
		for uint(router.retainer.Len()) > router.config.RetainCount {
			router.retainer.Remove(router.retainer.Front())
		}
	}

	return err
}

// Inbound returns the channel which transmits incoming data.
func (router *Router) Inbound() <-chan cemi.Message {
	return router.inbound
}

// Close closes the underlying socket and terminates the Router thereby.
func (router *Router) Close() {
	router.sock.Close()
}
<commit_msg>Wait when router indicates that he is busy<commit_after>package knx

import (
	"container/list"

	"sync"

	"time"

	"github.com/vapourismo/knx-go/knx/cemi"
	"github.com/vapourismo/knx-go/knx/proto"
)

// A RouterConfig determines certain properties of a Router.
type RouterConfig struct {
	// Specify how many sent messages to retain. This is important for when a router indicates that
	// it has lost some messages. If you do not expect to saturate the router, keep this low.
	RetainCount uint
}

// Default configuration elements
var (
	defaultRetainCount uint = 32

	DefaultRouterConfig = RouterConfig{
		defaultRetainCount,
	}
)

// checkRouterConfig validates the given RouterConfig.
func checkRouterConfig(config RouterConfig) RouterConfig {
	if config.RetainCount == 0 {
		config.RetainCount = defaultRetainCount
	}

	return config
}

// tryPushInbound sends the message through the channel. If the sending blocks, it will launch a
// goroutine which will do the sending.
func tryPushInbound(msg cemi.Message, inbound chan<- cemi.Message) {
	select {
	case inbound <- msg:

	default:
		go func() {
			// Since this goroutine decouples from the server goroutine, it might try to send when
			// the server closed the inbound channel. Sending to a closed channel will panic. But we
			// don't care, because cool guys don't look at explosions.
			defer func() { recover() }()
			inbound <- msg
		}()
	}
}

// A Router is a participant in a KNXnet/IP multicast group.
type Router struct {
	sock     Socket
	config   RouterConfig
	inbound  chan cemi.Message
	sendMu   sync.Mutex
	retainer *list.List
}

// serve listens for incoming routing-related packets.
func (router *Router) serve() {
	defer close(router.inbound)

	for msg := range router.sock.Inbound() {
		switch msg := msg.(type) {
		case *proto.RoutingInd:
			// Try to push it to the client without blocking this goroutine to long.
			tryPushInbound(msg.Payload, router.inbound)

		case *proto.RoutingBusy:
			router.sendMu.Lock()
			time.AfterFunc(msg.WaitTime, router.sendMu.Unlock)

		case *proto.RoutingLost:
			// TODO: Resend the last msg.Count frames.
		}
	}
}

// NewRouter creates a new Router that joins the given multicast group.
func NewRouter(multicastAddress string, config RouterConfig) (*Router, error) {
	sock, err := NewRoutingSocket(multicastAddress)
	if err != nil {
		return nil, err
	}

	r := &Router{
		sock:     sock,
		config:   checkRouterConfig(config),
		inbound:  make(chan cemi.Message),
		retainer: list.New(),
	}

	go r.serve()

	return r, nil
}

// Send transmits a packet.
func (router *Router) Send(data cemi.Message) error {
	// We lock this before doing any sending so the server goroutine can adjust the flow control.
	router.sendMu.Lock()
	defer router.sendMu.Unlock()

	err := router.sock.Send(&proto.RoutingInd{Payload: data})

	if err == nil {
		// Store this for potential resending.
		router.retainer.PushBack(data)

		// We don't want to keep more messages than necessary. The overhead needs to be removed.
		for uint(router.retainer.Len()) > router.config.RetainCount {
			router.retainer.Remove(router.retainer.Front())
		}
	}

	return err
}

// Inbound returns the channel which transmits incoming data.
func (router *Router) Inbound() <-chan cemi.Message {
	return router.inbound
}

// Close closes the underlying socket and terminates the Router thereby.
func (router *Router) Close() {
	router.sock.Close()
}
<|endoftext|>


<commit_before>package knx

import (
	"bytes"
	"net"
	"time"
	"github.com/vapourismo/knx-go/knx/proto"
)

// A Socket is a socket, duh.
type Socket interface {
	Send(payload proto.ServiceWriterTo) error
	Inbound() <-chan proto.Service
	Close() error
}

// NewClientSocket creates a new Socket which can used to exchange KNXnet/IP packets with a gateway.
func NewClientSocket(gatewayAddress string) (Socket, error) {
	addr, err := net.ResolveUDPAddr("udp4", gatewayAddress)
	if err != nil {
		return nil, err
	}

	conn, err := net.DialUDP("udp4", nil, addr)
	if err != nil {
		return nil, err
	}

	return makeUDPSocket(conn, addr), nil
}

// NewRoutingSocket creates a new Socket which can be used to exchange KNXnet/IP packets with a
// router.
func NewRoutingSocket(multicastAddress string) (Socket, error) {
	addr, err := net.ResolveUDPAddr("udp4", multicastAddress)
	if err != nil {
		return nil, err
	}

	conn, err := net.ListenMulticastUDP("udp4", nil, addr)
	if err != nil {
		return nil, err
	}

	return makeUDPSocket(conn, nil), nil
}

// UDP socket for KNXnet/IP packet exchange
type udpSocket struct {
	conn    *net.UDPConn
	inbound <-chan proto.Service
}

// makeUDPSocket configures the UDPConn and launches the receiver and sender workers.
func makeUDPSocket(conn *net.UDPConn, addr *net.UDPAddr) *udpSocket {
	conn.SetDeadline(time.Time{})

	inbound := make(chan proto.Service)
	go udpSocketReceiver(conn, addr, inbound)

	return &udpSocket{conn, inbound}
}

// Send transmits a KNXnet/IP packet.
func (sock *udpSocket) Send(payload proto.ServiceWriterTo) error {
	buffer := bytes.Buffer{}

	// Packet serialization
	_, err := proto.Pack(&buffer, payload)
	if err != nil {
		return err
	}

	log(sock.conn, "udpSocket", "<- %T %+v", payload, payload)

	// Transmission of the buffer contents
	_, err = sock.conn.Write(buffer.Bytes())
	if err != nil {
		return err
	}

	return nil
}

// Inbound provides a channel from which you can retrieve incoming packets.
func (sock *udpSocket) Inbound() <-chan proto.Service {
	return sock.inbound
}

// Close shuts the socket down. This will indirectly terminate the associated workers.
func (sock *udpSocket) Close() error {
	return sock.conn.Close()
}

// udpSocketReceiver is the receiver worker for udpSocket.
func udpSocketReceiver(conn *net.UDPConn, addr *net.UDPAddr, inbound chan<- proto.Service) {
	log(conn, "udpSocket", "Started receiver")
	defer log(conn, "udpSocket", "Stopped receiver")

	defer close(inbound)

	buffer := [1024]byte{}
	reader := bytes.NewReader(nil)

	for {
		len, sender, err := conn.ReadFromUDP(buffer[:])
		if err != nil {
			log(conn, "udpSocket", "Error during read: %v", err)
			return
		}

		// Validate sender origin if necessary
		if addr != nil && (!addr.IP.Equal(sender.IP) || addr.Port != sender.Port) {
			log(conn, "udpSocket", "Origin validation failed: %v (expected %v)", sender, addr)
			continue
		}

		reader.Reset(buffer[:len])

		var payload proto.Service
		_, err = proto.Unpack(reader, &payload)
		if err != nil {
			log(conn, "udpSocket", "Error during packet parsing: %v", err)
			continue
		}

		log(conn, "udpSocket", "-> %T %+v", payload, payload)

		inbound <- payload
	}
}
<commit_msg>Alloc new reader instead of resetting it every time<commit_after>package knx

import (
	"bytes"
	"net"
	"time"
	"github.com/vapourismo/knx-go/knx/proto"
)

// A Socket is a socket, duh.
type Socket interface {
	Send(payload proto.ServiceWriterTo) error
	Inbound() <-chan proto.Service
	Close() error
}

// NewClientSocket creates a new Socket which can used to exchange KNXnet/IP packets with a gateway.
func NewClientSocket(gatewayAddress string) (Socket, error) {
	addr, err := net.ResolveUDPAddr("udp4", gatewayAddress)
	if err != nil {
		return nil, err
	}

	conn, err := net.DialUDP("udp4", nil, addr)
	if err != nil {
		return nil, err
	}

	return makeUDPSocket(conn, addr), nil
}

// NewRoutingSocket creates a new Socket which can be used to exchange KNXnet/IP packets with a
// router.
func NewRoutingSocket(multicastAddress string) (Socket, error) {
	addr, err := net.ResolveUDPAddr("udp4", multicastAddress)
	if err != nil {
		return nil, err
	}

	conn, err := net.ListenMulticastUDP("udp4", nil, addr)
	if err != nil {
		return nil, err
	}

	return makeUDPSocket(conn, nil), nil
}

// UDP socket for KNXnet/IP packet exchange
type udpSocket struct {
	conn    *net.UDPConn
	inbound <-chan proto.Service
}

// makeUDPSocket configures the UDPConn and launches the receiver and sender workers.
func makeUDPSocket(conn *net.UDPConn, addr *net.UDPAddr) *udpSocket {
	conn.SetDeadline(time.Time{})

	inbound := make(chan proto.Service)
	go udpSocketReceiver(conn, addr, inbound)

	return &udpSocket{conn, inbound}
}

// Send transmits a KNXnet/IP packet.
func (sock *udpSocket) Send(payload proto.ServiceWriterTo) error {
	buffer := bytes.Buffer{}

	// Packet serialization
	_, err := proto.Pack(&buffer, payload)
	if err != nil {
		return err
	}

	log(sock.conn, "udpSocket", "<- %T %+v", payload, payload)

	// Transmission of the buffer contents
	_, err = sock.conn.Write(buffer.Bytes())
	if err != nil {
		return err
	}

	return nil
}

// Inbound provides a channel from which you can retrieve incoming packets.
func (sock *udpSocket) Inbound() <-chan proto.Service {
	return sock.inbound
}

// Close shuts the socket down. This will indirectly terminate the associated workers.
func (sock *udpSocket) Close() error {
	return sock.conn.Close()
}

// udpSocketReceiver is the receiver worker for udpSocket.
func udpSocketReceiver(conn *net.UDPConn, addr *net.UDPAddr, inbound chan<- proto.Service) {
	log(conn, "udpSocket", "Started receiver")
	defer log(conn, "udpSocket", "Stopped receiver")
	defer close(inbound)

	buffer := [1024]byte{}

	for {
		len, sender, err := conn.ReadFromUDP(buffer[:])
		if err != nil {
			log(conn, "udpSocket", "Error during read: %v", err)
			return
		}

		// Validate sender origin if necessary
		if addr != nil && (!addr.IP.Equal(sender.IP) || addr.Port != sender.Port) {
			log(conn, "udpSocket", "Origin validation failed: %v (expected %v)", sender, addr)
			continue
		}

		var payload proto.Service
		_, err = proto.Unpack(bytes.NewReader(buffer[:len]), &payload)
		if err != nil {
			log(conn, "udpSocket", "Error during packet parsing: %v", err)
			continue
		}

		log(conn, "udpSocket", "-> %T %+v", payload, payload)

		inbound <- payload
	}
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package instance

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/ernestio/ernestaws"
	"github.com/ernestio/ernestaws/credentials"
)

var (
	// ErrDatacenterIDInvalid ...
	ErrDatacenterIDInvalid = errors.New("Datacenter VPC ID invalid")
	// ErrDatacenterRegionInvalid ...
	ErrDatacenterRegionInvalid = errors.New("Datacenter Region invalid")
	// ErrDatacenterCredentialsInvalid ...
	ErrDatacenterCredentialsInvalid = errors.New("Datacenter credentials invalid")
	// ErrInstanceAWSIDInvalid ...
	ErrInstanceAWSIDInvalid = errors.New("Instance aws id invalid")
	// ErrNetworkInvalid ...
	ErrNetworkInvalid = errors.New("Network invalid")
	// ErrInstanceNameInvalid ...
	ErrInstanceNameInvalid = errors.New("Instance name invalid")
	// ErrInstanceImageInvalid ...
	ErrInstanceImageInvalid = errors.New("Instance image invalid")
	// ErrInstanceTypeInvalid ...
	ErrInstanceTypeInvalid = errors.New("Instance type invalid")
)

// Volume stores ebs volume data
type Volume struct {
	Name        string `json:"name"`
	Device      string `json:"device"`
	VolumeAWSID string `json:"volume_aws_id"`
}

// Event stores the template data
type Event struct {
	UUID                  string   `json:"_uuid"`
	BatchID               string   `json:"_batch_id"`
	ProviderType          string   `json:"_type"`
	VPCID                 string   `json:"vpc_id"`
	DatacenterRegion      string   `json:"datacenter_region"`
	DatacenterAccessKey   string   `json:"datacenter_secret"`
	DatacenterAccessToken string   `json:"datacenter_token"`
	NetworkAWSID          string   `json:"network_aws_id"`
	NetworkIsPublic       bool     `json:"network_is_public"`
	SecurityGroupAWSIDs   []string `json:"security_group_aws_ids"`
	InstanceAWSID         string   `json:"instance_aws_id,omitempty"`
	Name                  string   `json:"name"`
	Image                 string   `json:"image"`
	InstanceType          string   `json:"instance_type"`
	IP                    string   `json:"ip"`
	KeyPair               string   `json:"key_pair"`
	UserData              string   `json:"user_data"`
	PublicIP              string   `json:"public_ip"`
	ElasticIP             string   `json:"elastic_ip"`
	ElasticIPAWSID        string   `json:"elastic_ip_aws_id"`
	AssignElasticIP       bool     `json:"assign_elastic_ip"`
	Volumes               []Volume `json:"volumes"`
	ErrorMessage          string   `json:"error,omitempty"`
	Subject               string   `json:"-"`
	Body                  []byte   `json:"-"`
	CryptoKey             string   `json:"-"`
}

// New : Constructor
func New(subject string, body []byte, cryptoKey string) ernestaws.Event {
	n := Event{Subject: subject, Body: body, CryptoKey: cryptoKey}

	return &n
}

// GetBody : Gets the body for this event
func (ev *Event) GetBody() []byte {
	var err error
	if ev.Body, err = json.Marshal(ev); err != nil {
		log.Println(err.Error())
	}
	return ev.Body
}

// GetSubject : Gets the subject for this event
func (ev *Event) GetSubject() string {
	return ev.Subject
}

// Process : starts processing the current message
func (ev *Event) Process() (err error) {
	if err := json.Unmarshal(ev.Body, &ev); err != nil {
		ev.Error(err)
		return err
	}

	if err := ev.Validate(); err != nil {
		ev.Error(err)
		return err
	}

	return nil
}

// Error : Will respond the current event with an error
func (ev *Event) Error(err error) {
	log.Printf("Error: %s", err.Error())
	ev.ErrorMessage = err.Error()

	ev.Body, err = json.Marshal(ev)
}

// Validate checks if all criteria are met
func (ev *Event) Validate() error {
	if ev.VPCID == "" {
		return ErrDatacenterIDInvalid
	}

	if ev.DatacenterRegion == "" {
		return ErrDatacenterRegionInvalid
	}

	if ev.DatacenterAccessKey == "" || ev.DatacenterAccessToken == "" {
		return ErrDatacenterCredentialsInvalid
	}

	if ev.Subject != "instance.create.aws" {
		if ev.InstanceAWSID == "" {
			return ErrInstanceAWSIDInvalid
		}
	}

	if ev.Subject != "instance.delete.aws" {
		if ev.NetworkAWSID == "" {
			return ErrNetworkInvalid
		}
	}

	if ev.Name == "" {
		return ErrInstanceNameInvalid
	}

	if ev.Image == "" {
		return ErrInstanceImageInvalid
	}

	if ev.InstanceType == "" {
		return ErrInstanceTypeInvalid
	}

	return nil
}

// Create : Creates a instance object on aws
func (ev *Event) Create() error {
	svc := ev.getEC2Client()

	req := ec2.RunInstancesInput{
		SubnetId:         aws.String(ev.NetworkAWSID),
		ImageId:          aws.String(ev.Image),
		InstanceType:     aws.String(ev.InstanceType),
		PrivateIpAddress: aws.String(ev.IP),
		KeyName:          aws.String(ev.KeyPair),
		MaxCount:         aws.Int64(1),
		MinCount:         aws.Int64(1),
	}

	for _, sg := range ev.SecurityGroupAWSIDs {
		req.SecurityGroupIds = append(req.SecurityGroupIds, aws.String(sg))
	}

	if ev.UserData != "" {
		data := ev.encodeUserData(ev.UserData)
		req.UserData = aws.String(data)
	}

	resp, err := svc.RunInstances(&req)
	if err != nil {
		return err
	}

	builtInstance := ec2.DescribeInstancesInput{
		InstanceIds: []*string{resp.Instances[0].InstanceId},
	}

	err = svc.WaitUntilInstanceRunning(&builtInstance)
	if err != nil {
		return err
	}

	if ev.AssignElasticIP {
		ev.ElasticIP, ev.ElasticIPAWSID, err = ev.assignElasticIP(svc, *resp.Instances[0].InstanceId)
		if err != nil {
			return err
		}
	}

	ev.InstanceAWSID = *resp.Instances[0].InstanceId

	instance, err := ev.getInstanceByID(resp.Instances[0].InstanceId)
	if err != nil {
		return err
	}

	if instance.PublicIpAddress != nil {
		ev.PublicIP = *instance.PublicIpAddress
	}

	return ev.attachVolumes()
}

// Update : Updates a instance object on aws
func (ev *Event) Update() error {
	svc := ev.getEC2Client()

	builtInstance := ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	okInstance := ec2.DescribeInstanceStatusInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	err := svc.WaitUntilInstanceStatusOk(&okInstance)
	if err != nil {
		return err
	}

	stopreq := ec2.StopInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	// power off the instance
	_, err = svc.StopInstances(&stopreq)
	if err != nil {
		return err
	}

	err = svc.WaitUntilInstanceStopped(&builtInstance)
	if err != nil {
		return err
	}

	// resize the instance
	req := ec2.ModifyInstanceAttributeInput{
		InstanceId: aws.String(ev.InstanceAWSID),
		InstanceType: &ec2.AttributeValue{
			Value: aws.String(ev.InstanceType),
		},
	}

	_, err = svc.ModifyInstanceAttribute(&req)
	if err != nil {
		return err
	}

	// update instance security groups
	req = ec2.ModifyInstanceAttributeInput{
		InstanceId: aws.String(ev.InstanceAWSID),
		Groups:     []*string{},
	}

	for _, sg := range ev.SecurityGroupAWSIDs {
		req.Groups = append(req.Groups, aws.String(sg))
	}

	_, err = svc.ModifyInstanceAttribute(&req)
	if err != nil {
		return err
	}

	// power the instance back on
	startreq := ec2.StartInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	_, err = svc.StartInstances(&startreq)
	if err != nil {
		return err
	}

	err = svc.WaitUntilInstanceRunning(&builtInstance)
	if err != nil {
		return err
	}

	instance, err := ev.getInstanceByID(&ev.InstanceAWSID)
	if err != nil {
		return err
	}

	if instance.PublicIpAddress != nil {
		ev.PublicIP = *instance.PublicIpAddress
	}

	return ev.attachVolumes()
}

// Delete : Deletes a instance object on aws
func (ev *Event) Delete() error {
	svc := ev.getEC2Client()

	req := ec2.TerminateInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	_, err := svc.TerminateInstances(&req)
	if err != nil {
		return err
	}

	termreq := ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	err = svc.WaitUntilInstanceTerminated(&termreq)
	if err != nil {
		return err
	}

	if ev.ElasticIP != "" {
		dreq := &ec2.DisassociateAddressInput{
			AssociationId: aws.String(ev.ElasticIPAWSID),
			PublicIp:      aws.String(ev.ElasticIP),
		}

		_, err = svc.DisassociateAddress(dreq)
	}

	return err
}

// Get : Gets a instance object on aws
func (ev *Event) Get() error {
	return errors.New(ev.Subject + " not supported")
}

func (ev *Event) getEC2Client() *ec2.EC2 {
	creds, _ := credentials.NewStaticCredentials(ev.DatacenterAccessKey, ev.DatacenterAccessToken, ev.CryptoKey)
	return ec2.New(session.New(), &aws.Config{
		Region:      aws.String(ev.DatacenterRegion),
		Credentials: creds,
	})
}

func (ev *Event) assignElasticIP(svc *ec2.EC2, instanceID string) (string, string, error) {
	// Create Elastic IP
	resp, err := svc.AllocateAddress(nil)
	if err != nil {
		return "", "", err
	}

	req := ec2.AssociateAddressInput{
		InstanceId:   aws.String(instanceID),
		AllocationId: resp.AllocationId,
	}
	_, err = svc.AssociateAddress(&req)
	if err != nil {
		return "", "", err
	}

	return *resp.PublicIp, *resp.AllocationId, nil
}

func (ev *Event) getInstanceByID(id *string) (*ec2.Instance, error) {
	svc := ev.getEC2Client()

	req := ec2.DescribeInstancesInput{
		InstanceIds: []*string{id},
	}

	resp, err := svc.DescribeInstances(&req)
	if err != nil {
		return nil, err
	}

	if len(resp.Reservations) != 1 {
		return nil, errors.New("Could not find any instance reservations")
	}

	if len(resp.Reservations[0].Instances) != 1 {
		return nil, errors.New("Could not find an instance with that ID")
	}

	return resp.Reservations[0].Instances[0], nil
}

func (ev *Event) encodeUserData(data string) string {
	return base64.StdEncoding.EncodeToString([]byte(data))
}

func (ev *Event) attachVolumes() error {
	svc := ev.getEC2Client()

	instance, err := ev.getInstanceByID(&ev.InstanceAWSID)
	if err != nil {
		return err
	}

	for _, bdm := range instance.BlockDeviceMappings {
		if hasBlockDevice(ev.Volumes, bdm) {
			continue
		}

		req := &ec2.DetachVolumeInput{
			InstanceId: aws.String(ev.InstanceAWSID),
			VolumeId:   bdm.Ebs.VolumeId,
		}

		_, err = svc.DetachVolume(req)
		if err != nil {
			return err
		}
	}

	for _, vol := range ev.Volumes {
		// check volume doesn't exist
		if hasVolumeAttached(instance.BlockDeviceMappings, vol) {
			continue
		}

		req := &ec2.AttachVolumeInput{
			Device:     aws.String(vol.Device),
			VolumeId:   aws.String(vol.VolumeAWSID),
			InstanceId: aws.String(ev.InstanceAWSID),
		}

		_, err = svc.AttachVolume(req)
		if err != nil {
			return err
		}
	}

	return nil
}

func hasVolumeAttached(bdms []*ec2.InstanceBlockDeviceMapping, vol Volume) bool {
	for _, bdm := range bdms {
		if *bdm.Ebs.VolumeId == vol.VolumeAWSID || *bdm.DeviceName == vol.Device {
			return true
		}
	}

	return false
}

func hasBlockDevice(volumes []Volume, bdm *ec2.InstanceBlockDeviceMapping) bool {
	for _, vol := range volumes {
		if vol.VolumeAWSID == *bdm.Ebs.VolumeId || vol.Device == *bdm.DeviceName {
			return true
		}
	}

	return false
}
<commit_msg>correcting removal of root disk on instance update<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package instance

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/ernestio/ernestaws"
	"github.com/ernestio/ernestaws/credentials"
)

var (
	// ErrDatacenterIDInvalid ...
	ErrDatacenterIDInvalid = errors.New("Datacenter VPC ID invalid")
	// ErrDatacenterRegionInvalid ...
	ErrDatacenterRegionInvalid = errors.New("Datacenter Region invalid")
	// ErrDatacenterCredentialsInvalid ...
	ErrDatacenterCredentialsInvalid = errors.New("Datacenter credentials invalid")
	// ErrInstanceAWSIDInvalid ...
	ErrInstanceAWSIDInvalid = errors.New("Instance aws id invalid")
	// ErrNetworkInvalid ...
	ErrNetworkInvalid = errors.New("Network invalid")
	// ErrInstanceNameInvalid ...
	ErrInstanceNameInvalid = errors.New("Instance name invalid")
	// ErrInstanceImageInvalid ...
	ErrInstanceImageInvalid = errors.New("Instance image invalid")
	// ErrInstanceTypeInvalid ...
	ErrInstanceTypeInvalid = errors.New("Instance type invalid")
)

// Volume stores ebs volume data
type Volume struct {
	Name        string `json:"name"`
	Device      string `json:"device"`
	VolumeAWSID string `json:"volume_aws_id"`
}

// Event stores the template data
type Event struct {
	UUID                  string   `json:"_uuid"`
	BatchID               string   `json:"_batch_id"`
	ProviderType          string   `json:"_type"`
	VPCID                 string   `json:"vpc_id"`
	DatacenterRegion      string   `json:"datacenter_region"`
	DatacenterAccessKey   string   `json:"datacenter_secret"`
	DatacenterAccessToken string   `json:"datacenter_token"`
	NetworkAWSID          string   `json:"network_aws_id"`
	NetworkIsPublic       bool     `json:"network_is_public"`
	SecurityGroupAWSIDs   []string `json:"security_group_aws_ids"`
	InstanceAWSID         string   `json:"instance_aws_id,omitempty"`
	Name                  string   `json:"name"`
	Image                 string   `json:"image"`
	InstanceType          string   `json:"instance_type"`
	IP                    string   `json:"ip"`
	KeyPair               string   `json:"key_pair"`
	UserData              string   `json:"user_data"`
	PublicIP              string   `json:"public_ip"`
	ElasticIP             string   `json:"elastic_ip"`
	ElasticIPAWSID        string   `json:"elastic_ip_aws_id"`
	AssignElasticIP       bool     `json:"assign_elastic_ip"`
	Volumes               []Volume `json:"volumes"`
	ErrorMessage          string   `json:"error,omitempty"`
	Subject               string   `json:"-"`
	Body                  []byte   `json:"-"`
	CryptoKey             string   `json:"-"`
}

// New : Constructor
func New(subject string, body []byte, cryptoKey string) ernestaws.Event {
	n := Event{Subject: subject, Body: body, CryptoKey: cryptoKey}

	return &n
}

// GetBody : Gets the body for this event
func (ev *Event) GetBody() []byte {
	var err error
	if ev.Body, err = json.Marshal(ev); err != nil {
		log.Println(err.Error())
	}
	return ev.Body
}

// GetSubject : Gets the subject for this event
func (ev *Event) GetSubject() string {
	return ev.Subject
}

// Process : starts processing the current message
func (ev *Event) Process() (err error) {
	if err := json.Unmarshal(ev.Body, &ev); err != nil {
		ev.Error(err)
		return err
	}

	if err := ev.Validate(); err != nil {
		ev.Error(err)
		return err
	}

	return nil
}

// Error : Will respond the current event with an error
func (ev *Event) Error(err error) {
	log.Printf("Error: %s", err.Error())
	ev.ErrorMessage = err.Error()

	ev.Body, err = json.Marshal(ev)
}

// Validate checks if all criteria are met
func (ev *Event) Validate() error {
	if ev.VPCID == "" {
		return ErrDatacenterIDInvalid
	}

	if ev.DatacenterRegion == "" {
		return ErrDatacenterRegionInvalid
	}

	if ev.DatacenterAccessKey == "" || ev.DatacenterAccessToken == "" {
		return ErrDatacenterCredentialsInvalid
	}

	if ev.Subject != "instance.create.aws" {
		if ev.InstanceAWSID == "" {
			return ErrInstanceAWSIDInvalid
		}
	}

	if ev.Subject != "instance.delete.aws" {
		if ev.NetworkAWSID == "" {
			return ErrNetworkInvalid
		}
	}

	if ev.Name == "" {
		return ErrInstanceNameInvalid
	}

	if ev.Image == "" {
		return ErrInstanceImageInvalid
	}

	if ev.InstanceType == "" {
		return ErrInstanceTypeInvalid
	}

	return nil
}

// Create : Creates a instance object on aws
func (ev *Event) Create() error {
	svc := ev.getEC2Client()

	req := ec2.RunInstancesInput{
		SubnetId:         aws.String(ev.NetworkAWSID),
		ImageId:          aws.String(ev.Image),
		InstanceType:     aws.String(ev.InstanceType),
		PrivateIpAddress: aws.String(ev.IP),
		KeyName:          aws.String(ev.KeyPair),
		MaxCount:         aws.Int64(1),
		MinCount:         aws.Int64(1),
	}

	for _, sg := range ev.SecurityGroupAWSIDs {
		req.SecurityGroupIds = append(req.SecurityGroupIds, aws.String(sg))
	}

	if ev.UserData != "" {
		data := ev.encodeUserData(ev.UserData)
		req.UserData = aws.String(data)
	}

	resp, err := svc.RunInstances(&req)
	if err != nil {
		return err
	}

	builtInstance := ec2.DescribeInstancesInput{
		InstanceIds: []*string{resp.Instances[0].InstanceId},
	}

	err = svc.WaitUntilInstanceRunning(&builtInstance)
	if err != nil {
		return err
	}

	if ev.AssignElasticIP {
		ev.ElasticIP, ev.ElasticIPAWSID, err = ev.assignElasticIP(svc, *resp.Instances[0].InstanceId)
		if err != nil {
			return err
		}
	}

	ev.InstanceAWSID = *resp.Instances[0].InstanceId

	instance, err := ev.getInstanceByID(resp.Instances[0].InstanceId)
	if err != nil {
		return err
	}

	if instance.PublicIpAddress != nil {
		ev.PublicIP = *instance.PublicIpAddress
	}

	return ev.attachVolumes()
}

// Update : Updates a instance object on aws
func (ev *Event) Update() error {
	svc := ev.getEC2Client()

	builtInstance := ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	okInstance := ec2.DescribeInstanceStatusInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	err := svc.WaitUntilInstanceStatusOk(&okInstance)
	if err != nil {
		return err
	}

	stopreq := ec2.StopInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	// power off the instance
	_, err = svc.StopInstances(&stopreq)
	if err != nil {
		return err
	}

	err = svc.WaitUntilInstanceStopped(&builtInstance)
	if err != nil {
		return err
	}

	// resize the instance
	req := ec2.ModifyInstanceAttributeInput{
		InstanceId: aws.String(ev.InstanceAWSID),
		InstanceType: &ec2.AttributeValue{
			Value: aws.String(ev.InstanceType),
		},
	}

	_, err = svc.ModifyInstanceAttribute(&req)
	if err != nil {
		return err
	}

	// update instance security groups
	req = ec2.ModifyInstanceAttributeInput{
		InstanceId: aws.String(ev.InstanceAWSID),
		Groups:     []*string{},
	}

	for _, sg := range ev.SecurityGroupAWSIDs {
		req.Groups = append(req.Groups, aws.String(sg))
	}

	_, err = svc.ModifyInstanceAttribute(&req)
	if err != nil {
		return err
	}

	err = ev.attachVolumes()
	if err != nil {
		return err
	}

	// power the instance back on
	startreq := ec2.StartInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	_, err = svc.StartInstances(&startreq)
	if err != nil {
		return err
	}

	err = svc.WaitUntilInstanceRunning(&builtInstance)
	if err != nil {
		return err
	}

	instance, err := ev.getInstanceByID(&ev.InstanceAWSID)
	if err != nil {
		return err
	}

	if instance.PublicIpAddress != nil {
		ev.PublicIP = *instance.PublicIpAddress
	}

	return nil
}

// Delete : Deletes a instance object on aws
func (ev *Event) Delete() error {
	svc := ev.getEC2Client()

	req := ec2.TerminateInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	_, err := svc.TerminateInstances(&req)
	if err != nil {
		return err
	}

	termreq := ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(ev.InstanceAWSID)},
	}

	err = svc.WaitUntilInstanceTerminated(&termreq)
	if err != nil {
		return err
	}

	if ev.ElasticIP != "" {
		dreq := &ec2.DisassociateAddressInput{
			AssociationId: aws.String(ev.ElasticIPAWSID),
			PublicIp:      aws.String(ev.ElasticIP),
		}

		_, err = svc.DisassociateAddress(dreq)
	}

	return err
}

// Get : Gets a instance object on aws
func (ev *Event) Get() error {
	return errors.New(ev.Subject + " not supported")
}

func (ev *Event) getEC2Client() *ec2.EC2 {
	creds, _ := credentials.NewStaticCredentials(ev.DatacenterAccessKey, ev.DatacenterAccessToken, ev.CryptoKey)
	return ec2.New(session.New(), &aws.Config{
		Region:      aws.String(ev.DatacenterRegion),
		Credentials: creds,
	})
}

func (ev *Event) assignElasticIP(svc *ec2.EC2, instanceID string) (string, string, error) {
	// Create Elastic IP
	resp, err := svc.AllocateAddress(nil)
	if err != nil {
		return "", "", err
	}

	req := ec2.AssociateAddressInput{
		InstanceId:   aws.String(instanceID),
		AllocationId: resp.AllocationId,
	}
	_, err = svc.AssociateAddress(&req)
	if err != nil {
		return "", "", err
	}

	return *resp.PublicIp, *resp.AllocationId, nil
}

func (ev *Event) getInstanceByID(id *string) (*ec2.Instance, error) {
	svc := ev.getEC2Client()

	req := ec2.DescribeInstancesInput{
		InstanceIds: []*string{id},
	}

	resp, err := svc.DescribeInstances(&req)
	if err != nil {
		return nil, err
	}

	if len(resp.Reservations) != 1 {
		return nil, errors.New("Could not find any instance reservations")
	}

	if len(resp.Reservations[0].Instances) != 1 {
		return nil, errors.New("Could not find an instance with that ID")
	}

	return resp.Reservations[0].Instances[0], nil
}

func (ev *Event) encodeUserData(data string) string {
	return base64.StdEncoding.EncodeToString([]byte(data))
}

func (ev *Event) attachVolumes() error {
	svc := ev.getEC2Client()

	instance, err := ev.getInstanceByID(&ev.InstanceAWSID)
	if err != nil {
		return err
	}

	for _, bdm := range instance.BlockDeviceMappings {
		if hasBlockDevice(ev.Volumes, bdm) || *bdm.DeviceName == *instance.RootDeviceName {
			continue
		}

		req := &ec2.DetachVolumeInput{
			InstanceId: aws.String(ev.InstanceAWSID),
			VolumeId:   bdm.Ebs.VolumeId,
		}

		_, err = svc.DetachVolume(req)
		if err != nil {
			return err
		}
	}

	for _, vol := range ev.Volumes {
		// check volume doesn't exist
		if hasVolumeAttached(instance.BlockDeviceMappings, vol) {
			continue
		}

		req := &ec2.AttachVolumeInput{
			Device:     aws.String(vol.Device),
			VolumeId:   aws.String(vol.VolumeAWSID),
			InstanceId: aws.String(ev.InstanceAWSID),
		}

		_, err = svc.AttachVolume(req)
		if err != nil {
			return err
		}
	}

	return nil
}

func hasVolumeAttached(bdms []*ec2.InstanceBlockDeviceMapping, vol Volume) bool {
	for _, bdm := range bdms {
		if *bdm.Ebs.VolumeId == vol.VolumeAWSID || *bdm.DeviceName == vol.Device {
			return true
		}
	}

	return false
}

func hasBlockDevice(volumes []Volume, bdm *ec2.InstanceBlockDeviceMapping) bool {
	for _, vol := range volumes {
		if vol.VolumeAWSID == *bdm.Ebs.VolumeId || vol.Device == *bdm.DeviceName {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"os"
	"strings"
	"time"
)

// benchmarks: ab -r -k -l -n 30000 -c 5000 [uri]
func root(w http.ResponseWriter, r *http.Request) {
	// request uri = "http://localhost/live/luztv-livestream.w8889.m3u8?id=0x449484abb&wid=0xbc677870"
	// r.URL.Path[1:] = "live/luztv-livestream.w8889.m3u8" <=> r.URL.RawQuery = "id=0x449484abb&wid=0xbc677870"
	path := r.URL.Path[1:] // live/luztv-livestream.m3u8
	resp := ""
	if strings.Contains(path, ".m3u8") { // .m3u8 playlists
		if strings.Contains(path, "-playlist.m3u8") { // 1st identifying playlist
			// path = live/luztv-livestream-playlist.m3u8
			// recover the player cookie "rawstream" => ident
			var id int64
			var bps int
			var rawstream, key string
			a := strings.Split(path, "/")
			if len(a) == 2 {
				rawstream = strings.Replace(a[1], "-playlist.m3u8", "", -1) // luztv-livestream
			} else {
				http.NotFound(w, r)
				return
			}
			// get the bandwidth
			val, ok := Bw_int.Load(rawstream)
			if ok {
				bps = val.(int)
			} else {
				bps = 1000000 // 1 Mbps by default to avoid empty playlist.m3u8
			}
			// get the player cookie
			cookie, err := r.Cookie(rawstream)
			if err != nil {
				// not player cookie, let's create one
				mu_ident.Lock()
				ident++
				id = ident
				mu_ident.Unlock()
				key = fmt.Sprintf("%d", id)
			} else {
				// we have the cookie so we have the ident of this player
				key = cookie.Value // this is the id in string form
			}
			resp = fmt.Sprintf("#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=%d\n%s.wid%s.m3u8\n", bps, rawstream, key)
			expiration := time.Now().Add(24 * time.Hour)
			newcookie := http.Cookie{Name: rawstream, Value: key, Expires: expiration}
			http.SetCookie(w, &newcookie)
			w.Header().Set("Cache-Control", "no-cache")
			w.Header().Set("Content-Type", "application/vnd.apple.mpegurl")
			w.Header().Set("Access-Control-Allow-Headers", "*")
			w.Header().Set("Access-Control-Expose-Headers", "*")
			w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
			w.Header().Set("Access-Control-Allow-Credentials", "true")
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Accept-Ranges", "bytes")
			fmt.Fprintf(w, "%s", resp)
			return
		} else if strings.Contains(path, ".wid") { // recursive identified playlist
			// live/luztv-livestream.wid45006.m3u8
			var id int64
			tr := strings.Split(path, "/")
			spl := strings.Split(tr[len(tr)-1], ".")
			if len(spl) == 3 { // we response the content of the original .m3u8 playlist and record on database the stats info
				// we have to watch the referer, if allowed for this rawstream  ["rawstream"] = "domain1.com;domain2.com"
				val, ok := Referer.Load(spl[0])
				if ok {
					// get the referrer from request and compare to the one in the map (url: http://www.w3.org/hypertext/DataSources/Overview.html)
					if !strings.Contains(val.(string), getdomain(r.Referer())) {
						http.NotFound(w, r)
						return
					}
				}
				fmt.Sscanf(spl[1], "wid%d", &id) // wid9876
				file := spl[0] + ".m3u8"         // rawstream = spl[0]
				fileinfo, err := os.Stat(rootdir + "old/" + file)
				if err != nil {
					http.NotFound(w, r)
					return
				} else {
					// open the old/file.m3u8 (forecast pre-caching mechanism for CDNs)
					fr, errn := os.Open(rootdir + "old/" + file)
					if errn != nil {
						http.Error(w, "Internal Server Error", 500)
						return
					}
					defer fr.Close()
					if numgo < 1000000 { // if there are more than 1M goroutines working, live stats will stop for a while
						go createstats(r, spl[0], id)
					}
					w.Header().Set("Cache-Control", "no-cache")
					w.Header().Set("Content-Type", "application/vnd.apple.mpegurl")
					w.Header().Set("Access-Control-Allow-Headers", "*")
					w.Header().Set("Access-Control-Expose-Headers", "*")
					w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
					w.Header().Set("Access-Control-Allow-Credentials", "true")
					w.Header().Set("Access-Control-Allow-Origin", "*")
					w.Header().Set("Content-Length", fmt.Sprintf("%d", fileinfo.Size()))
					w.Header().Set("Accept-Ranges", "bytes")
					io.Copy(w, fr)
					return
				}
			} else {
				http.NotFound(w, r)
				return
			}
		} else {
			http.NotFound(w, r)
			return
		}
	} else if strings.Contains(path, ".ts") { // .TS segments
		// live/segment-56.ts
		fileinfo, err := os.Stat(rootdir + path)
		if err != nil {
			http.NotFound(w, r)
			return
		} else {
			fr, err := os.Open(rootdir + path)
			if err != nil {
				http.Error(w, "Internal Server Error", 500)
				return
			}
			defer fr.Close()
			w.Header().Set("Cache-Control", "max-age=300")
			w.Header().Set("Content-Type", "video/MP2T")
			w.Header().Set("Access-Control-Allow-Headers", "*")
			w.Header().Set("Access-Control-Expose-Headers", "*")
			w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
			w.Header().Set("Access-Control-Allow-Credentials", "true")
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Content-Length", fmt.Sprintf("%d", fileinfo.Size()))
			w.Header().Set("Accept-Ranges", "bytes")
			io.Copy(w, fr)
			return
		}
	} else if strings.Contains(path, ".lst") { // latest segment to pre-cache (forecaster map)
		// http://hlserver3/live/luztv-livestream.lst
		// tail -1 /var/segments/new/luztv-livestream.m3u8
		preload := false
		tr := strings.Split(path, "/")
		spl := strings.Split(tr[len(tr)-1], ".") // spl[0] = rawstream
		ip := getip(r.RemoteAddr)                // take near_proxy ip or r.Header["X-Cdn-Pop"] = [gsw] ???
		key := ip + "=" + spl[0]
		val, ok := Forecaster.Load(key)
		if ok {
			t := time.Now().Unix()
			if t-val.(int64) >= 5 { // if more than 5 secs
				preload = true
				Forecaster.Store(key, time.Now().Unix())
			}
		} else {
			preload = true
			Forecaster.Store(key, time.Now().Unix())
		}
		if preload {
			resp = getlatestseg(spl[0])
		}
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Access-Control-Allow-Headers", "*")
		w.Header().Set("Access-Control-Expose-Headers", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
		w.Header().Set("Access-Control-Allow-Credentials", "true")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Accept-Ranges", "bytes")
		fmt.Fprintf(w, "%s", resp) // here the response is the .ts file or just a blank string
	} else { // regular web content: http://hlserver.com/path_to_file
		file := rootdir + path
		fileinfo, err := os.Stat(file)
		if err != nil { // does not exist the path (file nor dir)
			http.NotFound(w, r)
			return
		} else if fileinfo.IsDir() { // it is a dir
			if strings.HasSuffix(file, "/") { // add /index.html to the end
				file = file + first_page + ".html"
			} else {
				file = file + "/" + first_page + ".html"
			}
		}
		fr, err := os.Open(file)
		if err != nil {
			http.Error(w, "Internal Server Error", 500)
			return
		}
		defer fr.Close()
		if session {
			if strings.Contains(r.URL.String(), "?err") {
				// replace  with an error text to show
				buf, _ := ioutil.ReadAll(fr)
				html := string(buf)
				html = strings.Replace(html, spanHTMLlogerr, ErrorText, -1)
				w.Header().Set("Cache-Control", "no-cache")
				w.Header().Set("Content-Type", mime.TypeByExtension(".html"))
				fmt.Fprint(w, html)
			} else {
				// Get the cookies
				filepart := strings.Split(file, ".")
				if (filepart[1] != "html") || (filepart[0] == (rootdir + first_page)) {
					w.Header().Set("Cache-Control", "no-cache")
					http.ServeContent(w, r, file, fileinfo.ModTime(), fr)
				} else {
					cookie, err := r.Cookie(CookieName)
					if err != nil {
						Error.Println("Cookie not found in the browser")
						http.Redirect(w, r, "/"+first_page+".html", http.StatusFound)
					} else {
						key := cookie.Value
						mu_user.RLock()
						_, ok := user_[key]
						mu_user.RUnlock()
						if ok {
							cookie.Expires = time.Now().Add(time.Duration(session_timeout) * time.Second)
							http.SetCookie(w, cookie)
							mu_user.Lock()
							time_[cookie.Value] = cookie.Expires
							mu_user.Unlock()
							w.Header().Set("Cache-Control", "no-cache")
							http.ServeContent(w, r, file, fileinfo.ModTime(), fr)
						} else {
							Error.Println("Cookie not found in the server")
							http.Redirect(w, r, "/"+first_page+".html", http.StatusFound)
						}
					}
				}
			}
		} else {
			w.Header().Set("Cache-Control", "no-cache")
			http.ServeContent(w, r, file, fileinfo.ModTime(), fr)
		}
	}
}

// just record in live.db @ table players (insert or update)
func createstats(r *http.Request, rawstream string, id int64) {
	// lets collect all the info to save in live.db
	var ipclient, ipproxy, username, streamname, os string
	var country, isocode, city string

	tr := strings.Split(rawstream, "-")
	if len(tr) != 2 {
		return
	}
	username, streamname = tr[0], tr[1]
	value, ok := r.Header["Remote-Ip"]
	if !ok {
		ipclient = getip(r.RemoteAddr)
	} else {
		ipclient = value[0]
	}
	ipproxy = getip(r.RemoteAddr)
	os = getos(r.UserAgent())
	country, isocode, city = geoIP(ipclient)
	timestamp := time.Now().Unix()

	// store everything in live.db
	mu_dblive.Lock()
	_, err := dblive.Exec("INSERT INTO players (`id`, `username`, `streamname`, `os`, `ipproxy`, `ipclient`, `isocode`, `country`, `city`, `timestamp`, `time`, `kilobytes`, `total_time`) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
		id, username, streamname, os, ipproxy, ipclient, isocode, country, city, timestamp, 0, 0, 0)
	mu_dblive.Unlock()
	if err != nil {
		if strings.Contains(err.Error(), "constraint") { // UNIQUE constraint failed: players.id // this id existed
			// let's get: timestamp, time, kilobytes, total_time and bandwith(rawstream) // ["luztv-livestream"] = 3780000 bps
			bw := 0
			bandwidth, ok := Bw_int.Load(rawstream)
			if ok {
				bw = bandwidth.(int)
			}
			var timestamp_db, time_db, kilobytes_db, total_time_db int64
			err := dblive.QueryRow("SELECT timestamp, time, kilobytes, total_time FROM players WHERE id = ?", id).Scan(&timestamp_db, &time_db, &kilobytes_db, &total_time_db)
			if err != nil {
				Error.Println(err)
				return
			}
			seconds := timestamp - int64(timestamp_db)
			if seconds > 30 { // reconnected from a previous disconn
				mu_dblive.Lock()
				dblive.Exec("UPDATE players SET time = 0, timestamp = ? WHERE id = ?", timestamp, id)
				mu_dblive.Unlock()
			} else { // still connected
				KBs := kilobytes_db + (int64(bw) * seconds / 8192)
				mu_dblive.Lock()
				dblive.Exec("UPDATE players SET time = ?, total_time = ?, kilobytes = ?, timestamp = ? WHERE id = ?", time_db+seconds, total_time_db+seconds, KBs, timestamp, id)
				mu_dblive.Unlock()
			}
		} else {
			Error.Println(err)
		}
	}
}
<commit_msg>added forecaster.js<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"os"
	"strings"
	"time"
)

// benchmarks: ab -r -k -l -n 30000 -c 5000 [uri]
func root(w http.ResponseWriter, r *http.Request) {
	// request uri = "http://localhost/live/luztv-livestream.w8889.m3u8?id=0x449484abb&wid=0xbc677870"
	// r.URL.Path[1:] = "live/luztv-livestream.w8889.m3u8" <=> r.URL.RawQuery = "id=0x449484abb&wid=0xbc677870"
	path := r.URL.Path[1:] // live/luztv-livestream.m3u8
	resp := ""
	if strings.Contains(path, ".m3u8") { // .m3u8 playlists
		if strings.Contains(path, "-playlist.m3u8") { // 1st identifying playlist
			// path = live/luztv-livestream-playlist.m3u8
			// recover the player cookie "rawstream" => ident
			var id int64
			var bps int
			var rawstream, key string
			a := strings.Split(path, "/")
			if len(a) == 2 {
				rawstream = strings.Replace(a[1], "-playlist.m3u8", "", -1) // luztv-livestream
			} else {
				http.NotFound(w, r)
				return
			}
			// get the bandwidth
			val, ok := Bw_int.Load(rawstream)
			if ok {
				bps = val.(int)
			} else {
				bps = 1000000 // 1 Mbps by default to avoid empty playlist.m3u8
			}
			// get the player cookie
			cookie, err := r.Cookie(rawstream)
			if err != nil {
				// not player cookie, let's create one
				mu_ident.Lock()
				ident++
				id = ident
				mu_ident.Unlock()
				key = fmt.Sprintf("%d", id)
			} else {
				// we have the cookie so we have the ident of this player
				key = cookie.Value // this is the id in string form
			}
			resp = fmt.Sprintf("#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=%d\n%s.wid%s.m3u8\n", bps, rawstream, key)
			expiration := time.Now().Add(24 * time.Hour)
			newcookie := http.Cookie{Name: rawstream, Value: key, Expires: expiration}
			http.SetCookie(w, &newcookie)
			w.Header().Set("Cache-Control", "no-cache")
			w.Header().Set("Content-Type", "application/vnd.apple.mpegurl")
			w.Header().Set("Access-Control-Allow-Headers", "*")
			w.Header().Set("Access-Control-Expose-Headers", "*")
			w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
			w.Header().Set("Access-Control-Allow-Credentials", "true")
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Accept-Ranges", "bytes")
			fmt.Fprintf(w, "%s", resp)
			return
		} else if strings.Contains(path, ".wid") { // recursive identified playlist
			// live/luztv-livestream.wid45006.m3u8
			var id int64
			tr := strings.Split(path, "/")
			spl := strings.Split(tr[len(tr)-1], ".")
			if len(spl) == 3 { // we response the content of the original .m3u8 playlist and record on database the stats info
				// we have to watch the referer, if allowed for this rawstream  ["rawstream"] = "domain1.com;domain2.com"
				val, ok := Referer.Load(spl[0])
				if ok {
					// get the referrer from request and compare to the one in the map (url: http://www.w3.org/hypertext/DataSources/Overview.html)
					if !strings.Contains(val.(string), getdomain(r.Referer())) {
						http.NotFound(w, r)
						return
					}
				}
				fmt.Sscanf(spl[1], "wid%d", &id) // wid9876
				file := spl[0] + ".m3u8"         // rawstream = spl[0]
				fileinfo, err := os.Stat(rootdir + "old/" + file)
				if err != nil {
					http.NotFound(w, r)
					return
				} else {
					// open the old/file.m3u8 (forecast pre-caching mechanism for CDNs)
					fr, errn := os.Open(rootdir + "old/" + file)
					if errn != nil {
						http.Error(w, "Internal Server Error", 500)
						return
					}
					defer fr.Close()
					if numgo < 1000000 { // if there are more than 1M goroutines working, live stats will stop for a while
						go createstats(r, spl[0], id)
					}
					w.Header().Set("Cache-Control", "no-cache")
					w.Header().Set("Content-Type", "application/vnd.apple.mpegurl")
					w.Header().Set("Access-Control-Allow-Headers", "*")
					w.Header().Set("Access-Control-Expose-Headers", "*")
					w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
					w.Header().Set("Access-Control-Allow-Credentials", "true")
					w.Header().Set("Access-Control-Allow-Origin", "*")
					w.Header().Set("Content-Length", fmt.Sprintf("%d", fileinfo.Size()))
					w.Header().Set("Accept-Ranges", "bytes")
					io.Copy(w, fr)
					return
				}
			} else {
				http.NotFound(w, r)
				return
			}
		} else {
			http.NotFound(w, r)
			return
		}
	} else if strings.Contains(path, ".ts") { // .TS segments
		// live/segment-56.ts
		fileinfo, err := os.Stat(rootdir + path)
		if err != nil {
			http.NotFound(w, r)
			return
		} else {
			fr, err := os.Open(rootdir + path)
			if err != nil {
				http.Error(w, "Internal Server Error", 500)
				return
			}
			defer fr.Close()
			w.Header().Set("Cache-Control", "max-age=300")
			w.Header().Set("Content-Type", "video/MP2T")
			w.Header().Set("Access-Control-Allow-Headers", "*")
			w.Header().Set("Access-Control-Expose-Headers", "*")
			w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
			w.Header().Set("Access-Control-Allow-Credentials", "true")
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Content-Length", fmt.Sprintf("%d", fileinfo.Size()))
			w.Header().Set("Accept-Ranges", "bytes")
			io.Copy(w, fr)
			return
		}
	} else if strings.Contains(path, "/forecaster.js") {
		// this code will send the javascript code for the forecaster mechanism of pre-caching
		// http://hlserver/rawstream/forecaster.js (path = rawstream/forecaster.js)
		tr := strings.Split(path, "/")                                                     // rawstream = tr[0]
		jscode := fmt.Sprintf("$.ajax({ url : \"/live/%s.lst\", type : 'HEAD' });", tr[0]) // the JS code is not complete, needs to reload every 10 seconds
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Access-Control-Allow-Headers", "*")
		w.Header().Set("Access-Control-Expose-Headers", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
		w.Header().Set("Access-Control-Allow-Credentials", "true")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Accept-Ranges", "bytes")
		fmt.Fprintf(w, "%s", jscode)
	} else if strings.Contains(path, ".lst") { // latest segment to pre-cache (forecaster map)
		// http://hlserver3/live/luztv-livestream.lst
		// tail -1 /var/segments/new/luztv-livestream.m3u8
		preload := false
		tr := strings.Split(path, "/")
		spl := strings.Split(tr[len(tr)-1], ".") // spl[0] = rawstream
		ip := getip(r.RemoteAddr)                // take near_proxy ip or r.Header["X-Cdn-Pop"] = [gsw] ???
		key := ip + "=" + spl[0]
		val, ok := Forecaster.Load(key)
		if ok {
			t := time.Now().Unix()
			if t-val.(int64) >= 5 { // if more than 5 secs
				preload = true
				Forecaster.Store(key, time.Now().Unix())
			}
		} else {
			preload = true
			Forecaster.Store(key, time.Now().Unix())
		}
		if preload {
			resp = getlatestseg(spl[0])
		}
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Access-Control-Allow-Headers", "*")
		w.Header().Set("Access-Control-Expose-Headers", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
		w.Header().Set("Access-Control-Allow-Credentials", "true")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Accept-Ranges", "bytes")
		fmt.Fprintf(w, "%s", resp) // here the response is the .ts file or just a blank string
	} else { // regular web content: http://hlserver.com/path_to_file
		file := rootdir + path
		fileinfo, err := os.Stat(file)
		if err != nil { // does not exist the path (file nor dir)
			http.NotFound(w, r)
			return
		} else if fileinfo.IsDir() { // it is a dir
			if strings.HasSuffix(file, "/") { // add /index.html to the end
				file = file + first_page + ".html"
			} else {
				file = file + "/" + first_page + ".html"
			}
		}
		fr, err := os.Open(file)
		if err != nil {
			http.Error(w, "Internal Server Error", 500)
			return
		}
		defer fr.Close()
		if session {
			if strings.Contains(r.URL.String(), "?err") {
				// replace  with an error text to show
				buf, _ := ioutil.ReadAll(fr)
				html := string(buf)
				html = strings.Replace(html, spanHTMLlogerr, ErrorText, -1)
				w.Header().Set("Cache-Control", "no-cache")
				w.Header().Set("Content-Type", mime.TypeByExtension(".html"))
				fmt.Fprint(w, html)
			} else {
				// Get the cookies
				filepart := strings.Split(file, ".")
				if (filepart[1] != "html") || (filepart[0] == (rootdir + first_page)) {
					w.Header().Set("Cache-Control", "no-cache")
					http.ServeContent(w, r, file, fileinfo.ModTime(), fr)
				} else {
					cookie, err := r.Cookie(CookieName)
					if err != nil {
						Error.Println("Cookie not found in the browser")
						http.Redirect(w, r, "/"+first_page+".html", http.StatusFound)
					} else {
						key := cookie.Value
						mu_user.RLock()
						_, ok := user_[key]
						mu_user.RUnlock()
						if ok {
							cookie.Expires = time.Now().Add(time.Duration(session_timeout) * time.Second)
							http.SetCookie(w, cookie)
							mu_user.Lock()
							time_[cookie.Value] = cookie.Expires
							mu_user.Unlock()
							w.Header().Set("Cache-Control", "no-cache")
							http.ServeContent(w, r, file, fileinfo.ModTime(), fr)
						} else {
							Error.Println("Cookie not found in the server")
							http.Redirect(w, r, "/"+first_page+".html", http.StatusFound)
						}
					}
				}
			}
		} else {
			w.Header().Set("Cache-Control", "no-cache")
			http.ServeContent(w, r, file, fileinfo.ModTime(), fr)
		}
	}
}

// just record in live.db @ table players (insert or update)
func createstats(r *http.Request, rawstream string, id int64) {
	// lets collect all the info to save in live.db
	var ipclient, ipproxy, username, streamname, os string
	var country, isocode, city string

	tr := strings.Split(rawstream, "-")
	if len(tr) != 2 {
		return
	}
	username, streamname = tr[0], tr[1]
	value, ok := r.Header["Remote-Ip"]
	if !ok {
		ipclient = getip(r.RemoteAddr)
	} else {
		ipclient = value[0]
	}
	ipproxy = getip(r.RemoteAddr)
	os = getos(r.UserAgent())
	country, isocode, city = geoIP(ipclient)
	timestamp := time.Now().Unix()

	// store everything in live.db
	mu_dblive.Lock()
	_, err := dblive.Exec("INSERT INTO players (`id`, `username`, `streamname`, `os`, `ipproxy`, `ipclient`, `isocode`, `country`, `city`, `timestamp`, `time`, `kilobytes`, `total_time`) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
		id, username, streamname, os, ipproxy, ipclient, isocode, country, city, timestamp, 0, 0, 0)
	mu_dblive.Unlock()
	if err != nil {
		if strings.Contains(err.Error(), "constraint") { // UNIQUE constraint failed: players.id // this id existed
			// let's get: timestamp, time, kilobytes, total_time and bandwith(rawstream) // ["luztv-livestream"] = 3780000 bps
			bw := 0
			bandwidth, ok := Bw_int.Load(rawstream)
			if ok {
				bw = bandwidth.(int)
			}
			var timestamp_db, time_db, kilobytes_db, total_time_db int64
			err := dblive.QueryRow("SELECT timestamp, time, kilobytes, total_time FROM players WHERE id = ?", id).Scan(&timestamp_db, &time_db, &kilobytes_db, &total_time_db)
			if err != nil {
				Error.Println(err)
				return
			}
			seconds := timestamp - int64(timestamp_db)
			if seconds > 30 { // reconnected from a previous disconn
				mu_dblive.Lock()
				dblive.Exec("UPDATE players SET time = 0, timestamp = ? WHERE id = ?", timestamp, id)
				mu_dblive.Unlock()
			} else { // still connected
				KBs := kilobytes_db + (int64(bw) * seconds / 8192)
				mu_dblive.Lock()
				dblive.Exec("UPDATE players SET time = ?, total_time = ?, kilobytes = ?, timestamp = ? WHERE id = ?", time_db+seconds, total_time_db+seconds, KBs, timestamp, id)
				mu_dblive.Unlock()
			}
		} else {
			Error.Println(err)
		}
	}
}
<|endoftext|>