<commit_before>// +build integration

// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 engine

import (
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/api"
	apicontainer "github.com/aws/amazon-ecs-agent/agent/api/container"
	"github.com/aws/amazon-ecs-agent/agent/api/container/status"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/stretchr/testify/assert"
)

const orderingTimeout = 90 * time.Second

// TestDependencyHealthCheck is a happy-case integration test that considers a workflow with a HEALTHY dependency
// condition. We ensure that the task can be both started and stopped.
func TestDependencyHealthCheck(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencyHealth"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"sleep 5 && exit 1"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "HEALTHY",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 60 && exit 0"}
	dependency.HealthCheckType = apicontainer.DockerHealthCheckType
	dependency.DockerConfig.Config = aws.String(alwaysHealthyHealthCheckConfig)

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// Both containers should start
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerRunningStateChange(t, taskEngine)
		verifyTaskIsRunning(stateChangeEvents, testTask)

		// Task should stop all at once
		verifyContainerStoppedStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencyComplete validates that the COMPLETE dependency condition will resolve when the child container exits
// with exit code 1. It ensures that the child is started and stopped before the parent starts.
func TestDependencyComplete(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencyComplete"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"sleep 5"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "COMPLETE",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 10 && exit 1"}
	dependency.Essential = false

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion and then exit
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// Second container starts after the first stops, task becomes running
		verifyContainerRunningStateChange(t, taskEngine)
		verifyTaskIsRunning(stateChangeEvents, testTask)

		// Last container stops and then the task stops
		verifyContainerStoppedStateChange(t, taskEngine)
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencySuccess validates that the SUCCESS dependency condition will resolve when the child container exits
// with exit code 0. It ensures that the child is started and stopped before the parent starts.
func TestDependencySuccess(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencySuccess"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "SUCCESS",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 10"}
	dependency.Essential = false

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// Second container starts after the first stops, task becomes running
		verifyContainerRunningStateChange(t, taskEngine)
		verifyTaskIsRunning(stateChangeEvents, testTask)

		// Last container stops and then the task stops
		verifyContainerStoppedStateChange(t, taskEngine)
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencySuccess validates that the SUCCESS dependency condition will fail when the child exits 1. This is a
// contrast to how COMPLETE behaves. Instead of starting the parent, the task should simply exit.
func TestDependencySuccessErrored(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencySuccessErrored"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "SUCCESS",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 10 && exit 1"}
	dependency.Essential = false

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// task should transition to stopped
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencySuccessTimeout
func TestDependencySuccessTimeout(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencySuccessTimeout"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "SUCCESS",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 15"}
	dependency.Essential = false

	// set the timeout to be shorter than the amount of time it takes to stop
	dependency.StartTimeout = 8

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// task should transition to stopped
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencyHealthyTimeout
func TestDependencyHealthyTimeout(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencyHealthyTimeout"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "HEALTHY",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 30"}
	dependency.HealthCheckType = apicontainer.DockerHealthCheckType

	// enter a healthcheck that will fail
	dependency.DockerConfig.Config = aws.String(`{
			"HealthCheck":{
				"Test":["CMD-SHELL", "exit 1"]
			}
		}`)

	// set the timeout. Duration doesn't matter since healthcheck will always be unhealthy.
	dependency.StartTimeout = 8

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// task should transition to stopped
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestShutdownOrder
func TestShutdownOrder(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testShutdownOrder"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	A := createTestContainerWithImageAndName(baseImageForOS, "parent")
	B := createTestContainerWithImageAndName(baseImageForOS, "parent")
	C := createTestContainerWithImageAndName(baseImageForOS, "parent")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.Essential = true
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "A",
			Condition:     "START",
		},
		{
			ContainerName: "B",
			Condition:     "START",
		},
		{
			ContainerName: "C",
			Condition:     "START",
		},
	}

	A.EntryPoint = &entryPointForOS
	A.Command = []string{"sleep 1000"}
	A.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "B",
			Condition:     "START",
		},
	}

	B.EntryPoint = &entryPointForOS
	B.Command = []string{"sleep 1000"}
	B.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "C",
			Condition:     "START",
		},
	}

	C.EntryPoint = &entryPointForOS
	C.Command = []string{"sleep 1000"}

	testTask.Containers = []*apicontainer.Container{
		parent,
		A,
		B,
		C,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// Everything should first progress to running
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerRunningStateChange(t, taskEngine)
		verifyTaskIsRunning(stateChangeEvents, testTask)

		// The shutdown order will now proceed. Parent will exit first since it has an explicit exit command.
		event := <-stateChangeEvents
		assert.Equal(t, event.(api.ContainerStateChange).Status, status.ContainerStopped)
		assert.Equal(t, event.(api.ContainerStateChange).ContainerName, "parent")

		// The dependency chain is A -> B -> C. We expect the inverse order to be followed for shutdown:
		// C shuts down, then B, then A
		expectedC := <-stateChangeEvents
		assert.Equal(t, expectedC.(api.ContainerStateChange).Status, status.ContainerStopped)
		assert.Equal(t, expectedC.(api.ContainerStateChange).ContainerName, "C")

		expectedB := <-stateChangeEvents
		assert.Equal(t, expectedB.(api.ContainerStateChange).Status, status.ContainerStopped)
		assert.Equal(t, expectedB.(api.ContainerStateChange).ContainerName, "B")

		expectedA := <-stateChangeEvents
		assert.Equal(t, expectedA.(api.ContainerStateChange).Status, status.ContainerStopped)
		assert.Equal(t, expectedA.(api.ContainerStateChange).ContainerName, "A")

		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

func waitFinished(t *testing.T, finished <-chan interface{}, duration time.Duration) {
	select {
	case <-finished:
		t.Log("Finished successfully.")
		return
	case <-time.After(90 * time.Second):
		t.Error("timed out after: ", duration)
		t.FailNow()
	}
}
<commit_msg>Fix Shutdown Order test<commit_after>// +build integration

// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 engine

import (
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/api"
	apicontainer "github.com/aws/amazon-ecs-agent/agent/api/container"
	"github.com/aws/amazon-ecs-agent/agent/api/container/status"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/stretchr/testify/assert"
)

const orderingTimeout = 90 * time.Second

// TestDependencyHealthCheck is a happy-case integration test that considers a workflow with a HEALTHY dependency
// condition. We ensure that the task can be both started and stopped.
func TestDependencyHealthCheck(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencyHealth"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"sleep 5 && exit 1"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "HEALTHY",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 60 && exit 0"}
	dependency.HealthCheckType = apicontainer.DockerHealthCheckType
	dependency.DockerConfig.Config = aws.String(alwaysHealthyHealthCheckConfig)

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// Both containers should start
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerRunningStateChange(t, taskEngine)
		verifyTaskIsRunning(stateChangeEvents, testTask)

		// Task should stop all at once
		verifyContainerStoppedStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencyComplete validates that the COMPLETE dependency condition will resolve when the child container exits
// with exit code 1. It ensures that the child is started and stopped before the parent starts.
func TestDependencyComplete(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencyComplete"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"sleep 5"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "COMPLETE",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 10 && exit 1"}
	dependency.Essential = false

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion and then exit
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// Second container starts after the first stops, task becomes running
		verifyContainerRunningStateChange(t, taskEngine)
		verifyTaskIsRunning(stateChangeEvents, testTask)

		// Last container stops and then the task stops
		verifyContainerStoppedStateChange(t, taskEngine)
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencySuccess validates that the SUCCESS dependency condition will resolve when the child container exits
// with exit code 0. It ensures that the child is started and stopped before the parent starts.
func TestDependencySuccess(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencySuccess"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "SUCCESS",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 10"}
	dependency.Essential = false

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// Second container starts after the first stops, task becomes running
		verifyContainerRunningStateChange(t, taskEngine)
		verifyTaskIsRunning(stateChangeEvents, testTask)

		// Last container stops and then the task stops
		verifyContainerStoppedStateChange(t, taskEngine)
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencySuccess validates that the SUCCESS dependency condition will fail when the child exits 1. This is a
// contrast to how COMPLETE behaves. Instead of starting the parent, the task should simply exit.
func TestDependencySuccessErrored(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencySuccessErrored"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "SUCCESS",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 10 && exit 1"}
	dependency.Essential = false

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// task should transition to stopped
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencySuccessTimeout
func TestDependencySuccessTimeout(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencySuccessTimeout"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "SUCCESS",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 15"}
	dependency.Essential = false

	// set the timeout to be shorter than the amount of time it takes to stop
	dependency.StartTimeout = 8

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// task should transition to stopped
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestDependencyHealthyTimeout
func TestDependencyHealthyTimeout(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testDependencyHealthyTimeout"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	dependency := createTestContainerWithImageAndName(baseImageForOS, "dependency")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"exit 0"}
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "dependency",
			Condition:     "HEALTHY",
		},
	}

	dependency.EntryPoint = &entryPointForOS
	dependency.Command = []string{"sleep 30"}
	dependency.HealthCheckType = apicontainer.DockerHealthCheckType

	// enter a healthcheck that will fail
	dependency.DockerConfig.Config = aws.String(`{
			"HealthCheck":{
				"Test":["CMD-SHELL", "exit 1"]
			}
		}`)

	// set the timeout. Duration doesn't matter since healthcheck will always be unhealthy.
	dependency.StartTimeout = 8

	testTask.Containers = []*apicontainer.Container{
		parent,
		dependency,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// First container should run to completion
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerStoppedStateChange(t, taskEngine)

		// task should transition to stopped
		verifyTaskIsStopped(stateChangeEvents, testTask)
		close(finished)
	}()

	waitFinished(t, finished, orderingTimeout)
}

// TestShutdownOrder
func TestShutdownOrder(t *testing.T) {
	shutdownOrderingTimeout := 120 * time.Second
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	taskArn := "testShutdownOrder"
	testTask := createTestTask(taskArn)

	parent := createTestContainerWithImageAndName(baseImageForOS, "parent")
	A := createTestContainerWithImageAndName(baseImageForOS, "A")
	B := createTestContainerWithImageAndName(baseImageForOS, "B")
	C := createTestContainerWithImageAndName(baseImageForOS, "C")

	parent.EntryPoint = &entryPointForOS
	parent.Command = []string{"echo hi"}
	parent.Essential = true
	parent.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "A",
			Condition:     "START",
		},
	}

	A.EntryPoint = &entryPointForOS
	A.Command = []string{"sleep 100"}
	A.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "B",
			Condition:     "START",
		},
	}

	B.EntryPoint = &entryPointForOS
	B.Command = []string{"sleep 100"}
	B.DependsOn = []apicontainer.DependsOn{
		{
			ContainerName: "C",
			Condition:     "START",
		},
	}

	C.EntryPoint = &entryPointForOS
	C.Command = []string{"sleep 100"}

	testTask.Containers = []*apicontainer.Container{
		parent,
		A,
		B,
		C,
	}

	go taskEngine.AddTask(testTask)

	finished := make(chan interface{})
	go func() {
		// Everything should first progress to running
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerRunningStateChange(t, taskEngine)
		verifyContainerRunningStateChange(t, taskEngine)
		verifyTaskIsRunning(stateChangeEvents, testTask)

		// The shutdown order will now proceed. Parent will exit first since it has an explicit exit command.
		event := <-stateChangeEvents
		assert.Equal(t, event.(api.ContainerStateChange).Status, status.ContainerStopped)
		assert.Equal(t, event.(api.ContainerStateChange).ContainerName, "parent")

		// The dependency chain is A -> B -> C. We expect the inverse order to be followed for shutdown:
		// A shuts down, then B, then C
		expectedC := <-stateChangeEvents
		assert.Equal(t, expectedC.(api.ContainerStateChange).Status, status.ContainerStopped)
		assert.Equal(t, expectedC.(api.ContainerStateChange).ContainerName, "A")

		expectedB := <-stateChangeEvents
		assert.Equal(t, expectedB.(api.ContainerStateChange).Status, status.ContainerStopped)
		assert.Equal(t, expectedB.(api.ContainerStateChange).ContainerName, "B")

		expectedA := <-stateChangeEvents
		assert.Equal(t, expectedA.(api.ContainerStateChange).Status, status.ContainerStopped)
		assert.Equal(t, expectedA.(api.ContainerStateChange).ContainerName, "C")

		close(finished)
	}()

	waitFinished(t, finished, shutdownOrderingTimeout)
}

func waitFinished(t *testing.T, finished <-chan interface{}, duration time.Duration) {
	select {
	case <-finished:
		t.Log("Finished successfully.")
		return
	case <-time.After(duration):
		t.Error("timed out after: ", duration)
		t.FailNow()
	}
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 statemanager implements simple constructs for saving and restoring
// state from disk.
// It provides the interface for a StateManager which can read/write arbitrary
// json data from/to disk.
package statemanager

import (
	"encoding/json"
	"errors"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/aws/amazon-ecs-agent/agent/logger"
)

const (
	// ECSDataVersion is the current version of saved data. Any backwards or
	// forwards incompatible changes to the data-format should increment this number
	// and retain the ability to read old data versions.
	// Version changes:
	// 1) initial
	// 2)
	//   a) Add 'ACSSeqNum' top level field (backwards compatible; technically
	//      forwards compatible but could cause resource constraint violations)
	//   b) remove 'DEAD', 'UNKNOWN' state from ever being marshalled (backward and
	//      forward compatible)
	// 3) Add 'Protocol' field to 'portMappings' and 'KnownPortBindings'
	// 4) Add 'DockerConfig' struct
	// 5) Add 'ImageStates' struct as part of ImageManager
	// 6)
	//   a) Refactor 'Internal' field in 'api.Container' to 'Type' enum
	//   b) Add 'ContainerResourcesProvisioned' as a new 'ContainerStatus' enum
	//   c) Add 'SteadyStateStatus' field to 'Container' struct
	//   d) Add 'ENIAttachments' struct
	//   e) Deprecate 'SteadyStateDependencies' in favor of 'TransitionDependencySet'
	//	 f) Add 'MetadataUpdated' field to 'api.Container'
	ECSDataVersion = 6

	// ecsDataFile specifies the filename in the ECS_DATADIR
	ecsDataFile = "ecs_agent_data.json"

	// minSaveInterval specifies how frequently to flush to disk
	minSaveInterval = 10 * time.Second
)

var log = logger.ForModule("statemanager")

// Saveable types should be able to be json serializable and deserializable
// Properly, this should have json.Marshaler/json.Unmarshaler here, but string
// and so on can be marshaled/unmarshaled sanely but don't fit those interfaces.
type Saveable interface{}

// Saver is a type that can be saved
type Saver interface {
	Save() error
	ForceSave() error
}

// Option functions are functions that may be used as part of constructing a new
// StateManager
type Option func(StateManager)

type saveableState map[string]*Saveable
type intermediateSaveableState map[string]json.RawMessage

// State is a struct of all data that should be saveable/loadable to disk. Each
// element should be json-serializable.
//
// Note, changing this to work with BinaryMarshaler or another more compact
// format would be fine, but everything already needs a json representation
// since that's our wire format and the extra space taken / IO-time is expected
// to be fairly negligable.
type state struct {
	Data saveableState

	Version int
}

type intermediateState struct {
	Data intermediateSaveableState
}

type versionOnlyState struct {
	Version int
}

type platformDependencies interface{}

// A StateManager can load and save state from disk.
// Load is not expected to return an error if there is no state to load.
type StateManager interface {
	Saver
	Load() error
}

type basicStateManager struct {
	statePath string // The path to a file in which state can be serialized

	state *state // pointers to the data we should save / load into

	saveTimesLock   sync.Mutex // guards save times
	lastSave        time.Time  //the last time a save completed
	nextPlannedSave time.Time  //the next time a save is planned

	savingLock sync.Mutex // guards marshal, write, move (on Linux), and load (on Windows)

	platformDependencies platformDependencies // platform-specific dependencies
}

// NewStateManager constructs a new StateManager which saves data at the
// location specified in cfg and operates under the given options.
// The returned StateManager will not save more often than every 10 seconds and
// will not reliably return errors with Save, but will log them appropriately.
func NewStateManager(cfg *config.Config, options ...Option) (StateManager, error) {
	fi, err := os.Stat(cfg.DataDir)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, errors.New("State manager DataDir must exist")
	}

	state := &state{
		Data:    make(saveableState),
		Version: ECSDataVersion,
	}
	manager := &basicStateManager{
		statePath: cfg.DataDir,
		state:     state,
	}

	for _, option := range options {
		option(manager)
	}

	manager.platformDependencies = newPlatformDependencies()

	return manager, nil
}

// AddSaveable is an option that adds a given saveable as one that should be saved
// under the given name. The name must be the same across uses of the
// statemanager (e.g. program invocations) for it to be serialized and
// deserialized correctly.
func AddSaveable(name string, saveable Saveable) Option {
	return (Option)(func(m StateManager) {
		manager, ok := m.(*basicStateManager)
		if !ok {
			log.Crit("Unable to add to state manager; unknown instantiation")
			return
		}
		manager.state.Data[name] = &saveable
	})
}

// Save triggers a save to file, though respects a minimum save interval to wait
// between saves.
func (manager *basicStateManager) Save() error {
	manager.saveTimesLock.Lock()
	defer manager.saveTimesLock.Unlock()
	if time.Since(manager.lastSave) >= minSaveInterval {
		// we can just save
		err := manager.ForceSave()
		manager.lastSave = time.Now()
		manager.nextPlannedSave = time.Time{} // re-zero it; assume all pending desires to save are fulfilled
		return err
	} else if manager.nextPlannedSave.IsZero() {
		// No save planned yet, we should plan one.
		next := manager.lastSave.Add(minSaveInterval)
		manager.nextPlannedSave = next
		go func() {
			time.Sleep(next.Sub(time.Now()))
			manager.Save()
		}()
	}
	// else nextPlannedSave wasn't Zero so there's a save planned elsewhere that'll
	// fulfill this
	return nil
}

// ForceSave saves the given State to a file. It is an atomic operation on POSIX
// systems (by Renaming over the target file).
// This function logs errors at will and does not necessarily expect the caller
// to handle the error because there's little a caller can do in general other
// than just keep going.
// In addition, the StateManager internally buffers save requests in order to
// only save at most every STATE_SAVE_INTERVAL.
func (manager *basicStateManager) ForceSave() error {
	manager.savingLock.Lock()
	defer manager.savingLock.Unlock()
	log.Info("Saving state!")
	s := manager.state
	s.Version = ECSDataVersion

	data, err := json.Marshal(s)
	if err != nil {
		log.Error("Error saving state; could not marshal data; this is odd", "err", err)
		return err
	}
	return manager.writeFile(data)
}

// Load reads state off the disk from the well-known filepath and loads it into
// the passed State object.
func (manager *basicStateManager) Load() error {
	s := manager.state
	log.Info("Loading state!")
	data, err := manager.readFile()
	if err != nil {
		log.Error("Error reading existing state file", "err", err)
		return err
	}
	if data == nil {
		return nil
	}
	// Dry-run to make sure this is a version we can understand
	err = manager.dryRun(data)
	if err != nil {
		return err
	}
	// Now load it into the actual state. The reason we do this with the
	// intermediate state is that we *must* unmarshal directly into the
	// "saveable" pointers we were given in AddSaveable; if we unmarshal
	// directly into a map with values of pointers, those pointers are lost.
	// We *must* unmarshal this way because the existing pointers could have
	// semi-initialized data (and are actually expected to)

	var intermediate intermediateState
	err = json.Unmarshal(data, &intermediate)
	if err != nil {
		log.Debug("Could not unmarshal into intermediate")
		return err
	}

	for key, rawJSON := range intermediate.Data {
		actualPointer, ok := manager.state.Data[key]
		if !ok {
			log.Error("Loading state: potentially malformed json key of " + key)
			continue
		}
		err = json.Unmarshal(rawJSON, actualPointer)
		if err != nil {
			log.Debug("Could not unmarshal into actual")
			return err
		}
	}

	log.Debug("Loaded state!", "state", s)
	return nil
}

func (manager *basicStateManager) dryRun(data []byte) error {
	// Dry-run to make sure this is a version we can understand
	tmps := versionOnlyState{}
	err := json.Unmarshal(data, &tmps)
	if err != nil {
		log.Crit("Could not unmarshal existing state; corrupted data?", "err", err, "data", data)
		return err
	}
	if tmps.Version > ECSDataVersion {
		strversion := strconv.Itoa(tmps.Version)
		return errors.New("Unsupported data format: Version " + strversion + " not " + strconv.Itoa(ECSDataVersion))
	}
	return nil
}
<commit_msg>statemanger: version bump to 7<commit_after>// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 statemanager implements simple constructs for saving and restoring
// state from disk.
// It provides the interface for a StateManager which can read/write arbitrary
// json data from/to disk.
package statemanager

import (
	"encoding/json"
	"errors"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/aws/amazon-ecs-agent/agent/logger"
)

const (
	// ECSDataVersion is the current version of saved data. Any backwards or
	// forwards incompatible changes to the data-format should increment this number
	// and retain the ability to read old data versions.
	// Version changes:
	// 1) initial
	// 2)
	//   a) Add 'ACSSeqNum' top level field (backwards compatible; technically
	//      forwards compatible but could cause resource constraint violations)
	//   b) remove 'DEAD', 'UNKNOWN' state from ever being marshalled (backward and
	//      forward compatible)
	// 3) Add 'Protocol' field to 'portMappings' and 'KnownPortBindings'
	// 4) Add 'DockerConfig' struct
	// 5) Add 'ImageStates' struct as part of ImageManager
	// 6)
	//   a) Refactor 'Internal' field in 'api.Container' to 'Type' enum
	//   b) Add 'ContainerResourcesProvisioned' as a new 'ContainerStatus' enum
	//   c) Add 'SteadyStateStatus' field to 'Container' struct
	//   d) Add 'ENIAttachments' struct
	//   e) Deprecate 'SteadyStateDependencies' in favor of 'TransitionDependencySet'
	// 7)
	//	 a) Add 'MetadataUpdated' field to 'api.Container'
	ECSDataVersion = 7

	// ecsDataFile specifies the filename in the ECS_DATADIR
	ecsDataFile = "ecs_agent_data.json"

	// minSaveInterval specifies how frequently to flush to disk
	minSaveInterval = 10 * time.Second
)

var log = logger.ForModule("statemanager")

// Saveable types should be able to be json serializable and deserializable
// Properly, this should have json.Marshaler/json.Unmarshaler here, but string
// and so on can be marshaled/unmarshaled sanely but don't fit those interfaces.
type Saveable interface{}

// Saver is a type that can be saved
type Saver interface {
	Save() error
	ForceSave() error
}

// Option functions are functions that may be used as part of constructing a new
// StateManager
type Option func(StateManager)

type saveableState map[string]*Saveable
type intermediateSaveableState map[string]json.RawMessage

// State is a struct of all data that should be saveable/loadable to disk. Each
// element should be json-serializable.
//
// Note, changing this to work with BinaryMarshaler or another more compact
// format would be fine, but everything already needs a json representation
// since that's our wire format and the extra space taken / IO-time is expected
// to be fairly negligable.
type state struct {
	Data saveableState

	Version int
}

type intermediateState struct {
	Data intermediateSaveableState
}

type versionOnlyState struct {
	Version int
}

type platformDependencies interface{}

// A StateManager can load and save state from disk.
// Load is not expected to return an error if there is no state to load.
type StateManager interface {
	Saver
	Load() error
}

type basicStateManager struct {
	statePath string // The path to a file in which state can be serialized

	state *state // pointers to the data we should save / load into

	saveTimesLock   sync.Mutex // guards save times
	lastSave        time.Time  //the last time a save completed
	nextPlannedSave time.Time  //the next time a save is planned

	savingLock sync.Mutex // guards marshal, write, move (on Linux), and load (on Windows)

	platformDependencies platformDependencies // platform-specific dependencies
}

// NewStateManager constructs a new StateManager which saves data at the
// location specified in cfg and operates under the given options.
// The returned StateManager will not save more often than every 10 seconds and
// will not reliably return errors with Save, but will log them appropriately.
func NewStateManager(cfg *config.Config, options ...Option) (StateManager, error) {
	fi, err := os.Stat(cfg.DataDir)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, errors.New("State manager DataDir must exist")
	}

	state := &state{
		Data:    make(saveableState),
		Version: ECSDataVersion,
	}
	manager := &basicStateManager{
		statePath: cfg.DataDir,
		state:     state,
	}

	for _, option := range options {
		option(manager)
	}

	manager.platformDependencies = newPlatformDependencies()

	return manager, nil
}

// AddSaveable is an option that adds a given saveable as one that should be saved
// under the given name. The name must be the same across uses of the
// statemanager (e.g. program invocations) for it to be serialized and
// deserialized correctly.
func AddSaveable(name string, saveable Saveable) Option {
	return (Option)(func(m StateManager) {
		manager, ok := m.(*basicStateManager)
		if !ok {
			log.Crit("Unable to add to state manager; unknown instantiation")
			return
		}
		manager.state.Data[name] = &saveable
	})
}

// Save triggers a save to file, though respects a minimum save interval to wait
// between saves.
func (manager *basicStateManager) Save() error {
	manager.saveTimesLock.Lock()
	defer manager.saveTimesLock.Unlock()
	if time.Since(manager.lastSave) >= minSaveInterval {
		// we can just save
		err := manager.ForceSave()
		manager.lastSave = time.Now()
		manager.nextPlannedSave = time.Time{} // re-zero it; assume all pending desires to save are fulfilled
		return err
	} else if manager.nextPlannedSave.IsZero() {
		// No save planned yet, we should plan one.
		next := manager.lastSave.Add(minSaveInterval)
		manager.nextPlannedSave = next
		go func() {
			time.Sleep(next.Sub(time.Now()))
			manager.Save()
		}()
	}
	// else nextPlannedSave wasn't Zero so there's a save planned elsewhere that'll
	// fulfill this
	return nil
}

// ForceSave saves the given State to a file. It is an atomic operation on POSIX
// systems (by Renaming over the target file).
// This function logs errors at will and does not necessarily expect the caller
// to handle the error because there's little a caller can do in general other
// than just keep going.
// In addition, the StateManager internally buffers save requests in order to
// only save at most every STATE_SAVE_INTERVAL.
func (manager *basicStateManager) ForceSave() error {
	manager.savingLock.Lock()
	defer manager.savingLock.Unlock()
	log.Info("Saving state!")
	s := manager.state
	s.Version = ECSDataVersion

	data, err := json.Marshal(s)
	if err != nil {
		log.Error("Error saving state; could not marshal data; this is odd", "err", err)
		return err
	}
	return manager.writeFile(data)
}

// Load reads state off the disk from the well-known filepath and loads it into
// the passed State object.
func (manager *basicStateManager) Load() error {
	s := manager.state
	log.Info("Loading state!")
	data, err := manager.readFile()
	if err != nil {
		log.Error("Error reading existing state file", "err", err)
		return err
	}
	if data == nil {
		return nil
	}
	// Dry-run to make sure this is a version we can understand
	err = manager.dryRun(data)
	if err != nil {
		return err
	}
	// Now load it into the actual state. The reason we do this with the
	// intermediate state is that we *must* unmarshal directly into the
	// "saveable" pointers we were given in AddSaveable; if we unmarshal
	// directly into a map with values of pointers, those pointers are lost.
	// We *must* unmarshal this way because the existing pointers could have
	// semi-initialized data (and are actually expected to)

	var intermediate intermediateState
	err = json.Unmarshal(data, &intermediate)
	if err != nil {
		log.Debug("Could not unmarshal into intermediate")
		return err
	}

	for key, rawJSON := range intermediate.Data {
		actualPointer, ok := manager.state.Data[key]
		if !ok {
			log.Error("Loading state: potentially malformed json key of " + key)
			continue
		}
		err = json.Unmarshal(rawJSON, actualPointer)
		if err != nil {
			log.Debug("Could not unmarshal into actual")
			return err
		}
	}

	log.Debug("Loaded state!", "state", s)
	return nil
}

func (manager *basicStateManager) dryRun(data []byte) error {
	// Dry-run to make sure this is a version we can understand
	tmps := versionOnlyState{}
	err := json.Unmarshal(data, &tmps)
	if err != nil {
		log.Crit("Could not unmarshal existing state; corrupted data?", "err", err, "data", data)
		return err
	}
	if tmps.Version > ECSDataVersion {
		strversion := strconv.Itoa(tmps.Version)
		return errors.New("Unsupported data format: Version " + strversion + " not " + strconv.Itoa(ECSDataVersion))
	}
	return nil
}
<|endoftext|>


<commit_before>// Package ping tries to ping a HTTP server through different ways
// Connection, Session (Head), Get and Post
package ping

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"regexp"
	"time"

	"github.com/mehrdadrad/mylg/cli"
)

// Ping represents HTTP ping request
type Ping struct {
	url     string
	host    string
	timeout time.Duration
	count   int
	method  string
	rAddr   net.Addr
	nsTime  time.Duration
	conn    net.Conn
}

type Result struct {
	StatusCode int
	ConnTime   float64
	TotalTime  float64
	Size       int
	Proto      string
	Server     string
	Status     string
}

// NewPing validate and constructs request object
func NewPing(URL string, timeout time.Duration) (*Ping, error) {
	URL, flag := cli.Flag(URL)
	// help
	if _, ok := flag["help"]; ok || URL == "" {
		help()
		return nil, fmt.Errorf("")
	}
	URL = Normalize(URL)
	u, err := url.Parse(URL)
	if err != nil {
		return &Ping{}, fmt.Errorf("cannot parse url")
	}
	sTime := time.Now()
	ipAddr, err := net.ResolveIPAddr("ip", u.Host)
	if err != nil {
		return &Ping{}, fmt.Errorf("cannot resolve %s: Unknown host", u.Host)
	}

	p := &Ping{
		url:     URL,
		host:    u.Host,
		rAddr:   ipAddr,
		nsTime:  time.Since(sTime),
		timeout: timeout,
	}
	// set count
	p.count = cli.SetFlag(flag, "c", 4).(int)
	// set method
	p.method = cli.SetFlag(flag, "m", "HEAD").(string)

	return p, nil
}

// Normalize
func Normalize(URL string) string {
	re := regexp.MustCompile(`(?i)https{0,1}://`)
	if !re.MatchString(URL) {
		URL = fmt.Sprintf("http://%s", URL)
	}
	return URL
}

// ping tries to create a TCP connection
func (p *Ping) pingConn() (Result, bool) {
	var (
		r     Result
		err   error
		sTime time.Time
	)

	sTime = time.Now()
	p.conn, err = net.DialTimeout("tcp", p.url, p.timeout*time.Second)
	r.ConnTime = time.Since(sTime).Seconds()
	if err != nil {
		print(err.Error())
		return r, false
	}

	p.rAddr = p.conn.RemoteAddr()
	p.conn.Close()
	return r, true
}

// pingHeadLoop tries number of connection
// with header information
func (p *Ping) pingHeadLoop() {
	pStrPrefix := "HTTP Response seq=%d, "
	pStrPostfix := "proto=%s, status=%d, time=%.3f ms\n"
	fmt.Printf("HPING %s (%s), Method: HEAD, DNSLookup: %.4f ms\n", p.host, p.rAddr, p.nsTime.Seconds()*1000)
	for i := 0; i < p.count; i++ {
		if r, ok := p.pingHead(); ok {
			fmt.Printf(pStrPrefix+pStrPostfix, i, r.Proto, r.StatusCode, r.TotalTime*1000)
		} else {
			fmt.Printf(pStrPrefix+"timeout\n", i)
		}
	}
}

// pingHeadLoop tries number of connection
// with header information
func (p *Ping) pingGetLoop() {
	pStrPrefix := "HTTP Response seq=%d, "
	pStrPostfix := "proto=%s, status=%d, size=%d Bytes, time=%.3f ms\n"
	fmt.Printf("HPING %s (%s), Method: GET, DNSLookup: %.4f ms\n", p.host, p.rAddr, p.nsTime.Seconds()*1000)
	for i := 0; i < p.count; i++ {
		if r, ok := p.pingGet(); ok {
			fmt.Printf(pStrPrefix+pStrPostfix, i, r.Proto, r.StatusCode, r.Size, r.TotalTime*1000)
		} else {
			fmt.Printf(pStrPrefix+"timeout\n", i)
		}
	}
}

// pingGet tries to ping a web server through http
func (p *Ping) pingGet() (Result, bool) {
	var (
		r     Result
		sTime time.Time
	)

	client := &http.Client{Timeout: 2 * time.Second}
	sTime = time.Now()
	resp, err := client.Get(p.url)
	if err != nil {
		return r, false
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	r.Size = len(body)
	r.TotalTime = time.Since(sTime).Seconds()
	if err != nil {
		return r, false
	}
	r.StatusCode = resp.StatusCode
	r.Proto = resp.Proto
	return r, true
}

// pingHead tries to ping a web server through http
func (p *Ping) pingHead() (Result, bool) {
	var (
		r     Result
		sTime time.Time
	)

	client := &http.Client{Timeout: 2 * time.Second}
	sTime = time.Now()
	resp, err := client.Head(p.url)
	if err != nil {
		return r, false
	}
	r.TotalTime = time.Since(sTime).Seconds()
	if err != nil {
		return r, false
	}
	r.StatusCode = resp.StatusCode
	r.Proto = resp.Proto
	return r, true
}

// pingNetHead tries to execute head command
func (p *Ping) pingNetHead() (Result, bool) {
	var (
		r     Result
		b     = make([]byte, 512)
		err   error
		sTime time.Time
	)

	sTime = time.Now()
	p.conn, err = net.DialTimeout("tcp", p.host+":80", p.timeout*time.Second)

	if err != nil {
		print(err.Error())
		return r, false
	}

	fmt.Fprintf(p.conn, "HEAD / HTTP/1.1\r\n\r\n")
	reader := bufio.NewReader(p.conn)
	n, _ := reader.Read(b)
	for key, regex := range map[string]string{"Proto": `(HTTP/\d\.\d)`, "Status": `HTTP/\d\.\d\s+(\d+)`, "Server": `server:\s+(.*)\n`} {
		re := regexp.MustCompile(regex)
		a := re.FindSubmatch(b[:n])
		if len(a) == 2 {
			f := reflect.ValueOf(&r).Elem().FieldByName(key)
			f.SetString(string(a[1]))
		}
	}
	r.TotalTime = time.Since(sTime).Seconds()

	p.rAddr = p.conn.RemoteAddr()
	p.conn.Close()
	return r, true

}

func (p *Ping) Run() {
	switch p.method {
	case "HEAD":
		p.pingHeadLoop()
	case "GET":
		p.pingGetLoop()
	}
}

func help() {
	println(`
    usage:
          hping [-c count][-m method] url
	`)
}
<commit_msg>fixed golint<commit_after>// Package ping tries to ping a HTTP server through different ways
// Connection, Session (Head), Get and Post
package ping

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"regexp"
	"time"

	"github.com/mehrdadrad/mylg/cli"
)

// Ping represents HTTP ping request
type Ping struct {
	url     string
	host    string
	timeout time.Duration
	count   int
	method  string
	rAddr   net.Addr
	nsTime  time.Duration
	conn    net.Conn
}

// Result holds Ping result
type Result struct {
	StatusCode int
	ConnTime   float64
	TotalTime  float64
	Size       int
	Proto      string
	Server     string
	Status     string
}

// NewPing validate and constructs request object
func NewPing(URL string, timeout time.Duration) (*Ping, error) {
	URL, flag := cli.Flag(URL)
	// help
	if _, ok := flag["help"]; ok || URL == "" {
		help()
		return nil, fmt.Errorf("")
	}
	URL = Normalize(URL)
	u, err := url.Parse(URL)
	if err != nil {
		return &Ping{}, fmt.Errorf("cannot parse url")
	}
	sTime := time.Now()
	ipAddr, err := net.ResolveIPAddr("ip", u.Host)
	if err != nil {
		return &Ping{}, fmt.Errorf("cannot resolve %s: Unknown host", u.Host)
	}

	p := &Ping{
		url:     URL,
		host:    u.Host,
		rAddr:   ipAddr,
		nsTime:  time.Since(sTime),
		timeout: timeout,
	}
	// set count
	p.count = cli.SetFlag(flag, "c", 4).(int)
	// set method
	p.method = cli.SetFlag(flag, "m", "HEAD").(string)

	return p, nil
}

// Normalize fixes scheme
func Normalize(URL string) string {
	re := regexp.MustCompile(`(?i)https{0,1}://`)
	if !re.MatchString(URL) {
		URL = fmt.Sprintf("http://%s", URL)
	}
	return URL
}

// ping tries to create a TCP connection
func (p *Ping) pingConn() (Result, bool) {
	var (
		r     Result
		err   error
		sTime time.Time
	)

	sTime = time.Now()
	p.conn, err = net.DialTimeout("tcp", p.url, p.timeout*time.Second)
	r.ConnTime = time.Since(sTime).Seconds()
	if err != nil {
		print(err.Error())
		return r, false
	}

	p.rAddr = p.conn.RemoteAddr()
	p.conn.Close()
	return r, true
}

// pingHeadLoop tries number of connection
// with header information
func (p *Ping) pingHeadLoop() {
	pStrPrefix := "HTTP Response seq=%d, "
	pStrPostfix := "proto=%s, status=%d, time=%.3f ms\n"
	fmt.Printf("HPING %s (%s), Method: HEAD, DNSLookup: %.4f ms\n", p.host, p.rAddr, p.nsTime.Seconds()*1000)
	for i := 0; i < p.count; i++ {
		if r, ok := p.pingHead(); ok {
			fmt.Printf(pStrPrefix+pStrPostfix, i, r.Proto, r.StatusCode, r.TotalTime*1000)
		} else {
			fmt.Printf(pStrPrefix+"timeout\n", i)
		}
	}
}

// pingHeadLoop tries number of connection
// with header information
func (p *Ping) pingGetLoop() {
	pStrPrefix := "HTTP Response seq=%d, "
	pStrPostfix := "proto=%s, status=%d, size=%d Bytes, time=%.3f ms\n"
	fmt.Printf("HPING %s (%s), Method: GET, DNSLookup: %.4f ms\n", p.host, p.rAddr, p.nsTime.Seconds()*1000)
	for i := 0; i < p.count; i++ {
		if r, ok := p.pingGet(); ok {
			fmt.Printf(pStrPrefix+pStrPostfix, i, r.Proto, r.StatusCode, r.Size, r.TotalTime*1000)
		} else {
			fmt.Printf(pStrPrefix+"timeout\n", i)
		}
	}
}

// pingGet tries to ping a web server through http
func (p *Ping) pingGet() (Result, bool) {
	var (
		r     Result
		sTime time.Time
	)

	client := &http.Client{Timeout: 2 * time.Second}
	sTime = time.Now()
	resp, err := client.Get(p.url)
	if err != nil {
		return r, false
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	r.Size = len(body)
	r.TotalTime = time.Since(sTime).Seconds()
	if err != nil {
		return r, false
	}
	r.StatusCode = resp.StatusCode
	r.Proto = resp.Proto
	return r, true
}

// pingHead tries to ping a web server through http
func (p *Ping) pingHead() (Result, bool) {
	var (
		r     Result
		sTime time.Time
	)

	client := &http.Client{Timeout: 2 * time.Second}
	sTime = time.Now()
	resp, err := client.Head(p.url)
	if err != nil {
		return r, false
	}
	r.TotalTime = time.Since(sTime).Seconds()
	if err != nil {
		return r, false
	}
	r.StatusCode = resp.StatusCode
	r.Proto = resp.Proto
	return r, true
}

// pingNetHead tries to execute head command
func (p *Ping) pingNetHead() (Result, bool) {
	var (
		r     Result
		b     = make([]byte, 512)
		err   error
		sTime time.Time
	)

	sTime = time.Now()
	p.conn, err = net.DialTimeout("tcp", p.host+":80", p.timeout*time.Second)

	if err != nil {
		print(err.Error())
		return r, false
	}

	fmt.Fprintf(p.conn, "HEAD / HTTP/1.1\r\n\r\n")
	reader := bufio.NewReader(p.conn)
	n, _ := reader.Read(b)
	for key, regex := range map[string]string{"Proto": `(HTTP/\d\.\d)`, "Status": `HTTP/\d\.\d\s+(\d+)`, "Server": `server:\s+(.*)\n`} {
		re := regexp.MustCompile(regex)
		a := re.FindSubmatch(b[:n])
		if len(a) == 2 {
			f := reflect.ValueOf(&r).Elem().FieldByName(key)
			f.SetString(string(a[1]))
		}
	}
	r.TotalTime = time.Since(sTime).Seconds()

	p.rAddr = p.conn.RemoteAddr()
	p.conn.Close()
	return r, true

}

// Run tries to run ping loop based on the method
func (p *Ping) Run() {
	switch p.method {
	case "HEAD":
		p.pingHeadLoop()
	case "GET":
		p.pingGetLoop()
	}
}

// help shows ping help
func help() {
	println(`
    usage:
          hping [-c count][-m method] url
	`)
}
<|endoftext|>


<commit_before>package search

import (
	iopt "github.com/algolia/algoliasearch-client-go/v3/algolia/internal/opt"
	"github.com/algolia/algoliasearch-client-go/v3/algolia/opt"
)

type IndexOperation struct {
	Destination string   `json:"destination"`
	Operation   string   `json:"operation"`
	Scopes      []string `json:"scope,omitempty"`
}

type deleteByReq struct {
	AroundLatLng      *opt.AroundLatLngOption      `json:"aroundLatLng,omitempty"`
	AroundRadius      *opt.AroundRadiusOption      `json:"aroundRadius,omitempty"`
	FacetFilters      *opt.FacetFiltersOption      `json:"facetFilters,omitempty"`
	Filters           *opt.FiltersOption           `json:"filters,omitempty"`
	InsideBoundingBox *opt.InsideBoundingBoxOption `json:"insideBoundingBox,omitempty"`
	InsidePolygon     *opt.InsidePolygonOption     `json:"insidePolygon,omitempty"`
	NumericFilters    *opt.NumericFiltersOption    `json:"numericFilters,omitempty"`
}

func newDeleteByReq(opts ...interface{}) deleteByReq {
	return deleteByReq{
		AroundLatLng:      iopt.ExtractAroundLatLng(opts...),
		AroundRadius:      iopt.ExtractAroundRadius(opts...),
		FacetFilters:      iopt.ExtractFacetFilters(opts...),
		Filters:           iopt.ExtractFilters(opts...),
		InsideBoundingBox: iopt.ExtractInsideBoundingBox(opts...),
		InsidePolygon:     iopt.ExtractInsidePolygon(opts...),
		NumericFilters:    iopt.ExtractNumericFilters(opts...),
	}
}

type getObjectsReq struct {
	IndexName            string                          `json:"indexName"`
	ObjectID             string                          `json:"objectID"`
	AttributesToRetrieve *opt.AttributesToRetrieveOption `json:"attributesToRetrieve,omitempty"`
}

type batchReq struct {
	Requests []BatchOperation `json:"requests,omitempty"`
}
<commit_msg>feat(search): expose built-in operations for partial updates with PartialUpdateOperation<commit_after>package search

import (
	iopt "github.com/algolia/algoliasearch-client-go/v3/algolia/internal/opt"
	"github.com/algolia/algoliasearch-client-go/v3/algolia/opt"
)

type IndexOperation struct {
	Destination string   `json:"destination"`
	Operation   string   `json:"operation"`
	Scopes      []string `json:"scope,omitempty"`
}

type deleteByReq struct {
	AroundLatLng      *opt.AroundLatLngOption      `json:"aroundLatLng,omitempty"`
	AroundRadius      *opt.AroundRadiusOption      `json:"aroundRadius,omitempty"`
	FacetFilters      *opt.FacetFiltersOption      `json:"facetFilters,omitempty"`
	Filters           *opt.FiltersOption           `json:"filters,omitempty"`
	InsideBoundingBox *opt.InsideBoundingBoxOption `json:"insideBoundingBox,omitempty"`
	InsidePolygon     *opt.InsidePolygonOption     `json:"insidePolygon,omitempty"`
	NumericFilters    *opt.NumericFiltersOption    `json:"numericFilters,omitempty"`
}

func newDeleteByReq(opts ...interface{}) deleteByReq {
	return deleteByReq{
		AroundLatLng:      iopt.ExtractAroundLatLng(opts...),
		AroundRadius:      iopt.ExtractAroundRadius(opts...),
		FacetFilters:      iopt.ExtractFacetFilters(opts...),
		Filters:           iopt.ExtractFilters(opts...),
		InsideBoundingBox: iopt.ExtractInsideBoundingBox(opts...),
		InsidePolygon:     iopt.ExtractInsidePolygon(opts...),
		NumericFilters:    iopt.ExtractNumericFilters(opts...),
	}
}

type getObjectsReq struct {
	IndexName            string                          `json:"indexName"`
	ObjectID             string                          `json:"objectID"`
	AttributesToRetrieve *opt.AttributesToRetrieveOption `json:"attributesToRetrieve,omitempty"`
}

type batchReq struct {
	Requests []BatchOperation `json:"requests,omitempty"`
}

type PartialUpdateOperation struct {
	Operation string      `json:"_operation"`
	Value     interface{} `json:"value"`
}
<|endoftext|>


<commit_before>package tables_test

import (
	"context"
	"database/sql"
	"fmt"
	"reflect"
	"testing"
	"time"

	"github.com/matrix-org/dendrite/internal/sqlutil"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/dendrite/test"
	"github.com/matrix-org/dendrite/userapi/api"
	"github.com/matrix-org/dendrite/userapi/storage/postgres"
	"github.com/matrix-org/dendrite/userapi/storage/sqlite3"
	"github.com/matrix-org/dendrite/userapi/storage/tables"
	"github.com/matrix-org/dendrite/userapi/types"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
)

func mustMakeDBs(t *testing.T, dbType test.DBType) (
	*sql.DB, tables.AccountsTable, tables.DevicesTable, tables.StatsTable, func(),
) {
	t.Helper()

	var (
		accTable   tables.AccountsTable
		devTable   tables.DevicesTable
		statsTable tables.StatsTable
		err        error
	)

	connStr, close := test.PrepareDBConnectionString(t, dbType)
	db, err := sqlutil.Open(&config.DatabaseOptions{
		ConnectionString: config.DataSource(connStr),
	}, nil)
	if err != nil {
		t.Fatalf("failed to open db: %s", err)
	}

	switch dbType {
	case test.DBTypeSQLite:
		accTable, err = sqlite3.NewSQLiteAccountsTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to create acc db: %v", err)
		}
		devTable, err = sqlite3.NewSQLiteDevicesTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to open device db: %v", err)
		}
		statsTable, err = sqlite3.NewSQLiteStatsTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to open stats db: %v", err)
		}
	case test.DBTypePostgres:
		accTable, err = postgres.NewPostgresAccountsTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to create acc db: %v", err)
		}
		devTable, err = postgres.NewPostgresDevicesTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to open device db: %v", err)
		}
		statsTable, err = postgres.NewPostgresStatsTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to open stats db: %v", err)
		}
	}

	return db, accTable, devTable, statsTable, close
}

func mustMakeAccountAndDevice(
	t *testing.T,
	ctx context.Context,
	accDB tables.AccountsTable,
	devDB tables.DevicesTable,
	localpart string,
	accType api.AccountType,
	userAgent string,
) {
	t.Helper()

	appServiceID := ""
	if accType == api.AccountTypeAppService {
		appServiceID = util.RandomString(16)
	}

	_, err := accDB.InsertAccount(ctx, nil, localpart, "", appServiceID, accType)
	if err != nil {
		t.Fatalf("unable to create account: %v", err)
	}
	_, err = devDB.InsertDevice(ctx, nil, "deviceID", localpart, util.RandomString(16), nil, "", userAgent)
	if err != nil {
		t.Fatalf("unable to create device: %v", err)
	}
}

func mustUpdateDeviceLastSeen(
	t *testing.T,
	ctx context.Context,
	db *sql.DB,
	localpart string,
	timestamp time.Time,
) {
	t.Helper()
	_, err := db.ExecContext(ctx, "UPDATE userapi_devices SET last_seen_ts = $1 WHERE localpart = $2", gomatrixserverlib.AsTimestamp(timestamp), localpart)
	if err != nil {
		t.Fatalf("unable to update device last seen")
	}
}

func mustUserUpdateRegistered(
	t *testing.T,
	ctx context.Context,
	db *sql.DB,
	localpart string,
	timestamp time.Time,
) {
	_, err := db.ExecContext(ctx, "UPDATE userapi_accounts SET created_ts = $1 WHERE localpart = $2", gomatrixserverlib.AsTimestamp(timestamp), localpart)
	if err != nil {
		t.Fatalf("unable to update device last seen")
	}
}

// These tests must run sequentially, as they build up on each other
func Test_UserStatistics(t *testing.T) {

	ctx := context.Background()
	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
		db, accDB, devDB, statsDB, close := mustMakeDBs(t, dbType)
		defer close()
		wantType := "SQLite"
		if dbType == test.DBTypePostgres {
			wantType = "Postgres"
		}

		t.Run(fmt.Sprintf("want %s database engine", wantType), func(t *testing.T) {
			_, gotDB, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			if wantType != gotDB.Engine { // can't use DeepEqual, as the Version might differ
				t.Errorf("UserStatistics() got DB engine = %+v, want %s", gotDB.Engine, wantType)
			}
		})

		t.Run("Want Users", func(t *testing.T) {
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user1", api.AccountTypeUser, "Element Android")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user2", api.AccountTypeUser, "Element iOS")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user3", api.AccountTypeUser, "Element web")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user4", api.AccountTypeGuest, "Element Electron")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user5", api.AccountTypeAdmin, "gecko")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user6", api.AccountTypeAppService, "gecko")
			gotStats, _, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			wantStats := &types.UserStatistics{
				RegisteredUsersByType: map[string]int64{
					"native":  4,
					"guest":   1,
					"bridged": 1,
				},
				R30Users: map[string]int64{},
				R30UsersV2: map[string]int64{
					"ios":      0,
					"android":  0,
					"web":      0,
					"electron": 0,
					"all":      0,
				},
				AllUsers:        6,
				NonBridgedUsers: 5,
				DailyUsers:      6,
				MonthlyUsers:    6,
			}
			if !reflect.DeepEqual(gotStats, wantStats) {
				t.Errorf("UserStatistics() gotStats = \n%+v\nwant\n%+v", gotStats, wantStats)
			}
		})

		t.Run("Users not active for one/two month", func(t *testing.T) {
			mustUpdateDeviceLastSeen(t, ctx, db, "user1", time.Now().AddDate(0, -2, 0))
			mustUpdateDeviceLastSeen(t, ctx, db, "user2", time.Now().AddDate(0, -1, 0))
			gotStats, _, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			wantStats := &types.UserStatistics{
				RegisteredUsersByType: map[string]int64{
					"native":  4,
					"guest":   1,
					"bridged": 1,
				},
				R30Users: map[string]int64{},
				R30UsersV2: map[string]int64{
					"ios":      0,
					"android":  0,
					"web":      0,
					"electron": 0,
					"all":      0,
				},
				AllUsers:        6,
				NonBridgedUsers: 5,
				DailyUsers:      4,
				MonthlyUsers:    4,
			}
			if !reflect.DeepEqual(gotStats, wantStats) {
				t.Errorf("UserStatistics() gotStats = \n%+v\nwant\n%+v", gotStats, wantStats)
			}
		})

		/* R30Users counts the number of 30 day retained users, defined as:
		- Users who have created their accounts more than 30 days ago
		- Where last seen at most 30 days ago
		- Where account creation and last_seen are > 30 days apart
		*/
		t.Run("R30Users tests", func(t *testing.T) {
			mustUserUpdateRegistered(t, ctx, db, "user1", time.Now().AddDate(0, -2, 0))
			mustUpdateDeviceLastSeen(t, ctx, db, "user1", time.Now())
			mustUserUpdateRegistered(t, ctx, db, "user4", time.Now().AddDate(0, -2, 0))
			mustUpdateDeviceLastSeen(t, ctx, db, "user4", time.Now())
			startTime := time.Now().AddDate(0, 0, -2)
			err := statsDB.UpdateUserDailyVisits(ctx, nil, startTime, startTime.Truncate(time.Hour*24).Add(time.Hour))
			if err != nil {
				t.Fatalf("unable to update daily visits stats: %v", err)
			}

			gotStats, _, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			wantStats := &types.UserStatistics{
				RegisteredUsersByType: map[string]int64{
					"native":  3,
					"bridged": 1,
				},
				R30Users: map[string]int64{
					"all":      2,
					"android":  1,
					"electron": 1,
				},
				R30UsersV2: map[string]int64{
					"ios":      0,
					"android":  0,
					"web":      0,
					"electron": 0,
					"all":      0,
				},
				AllUsers:        6,
				NonBridgedUsers: 5,
				DailyUsers:      5,
				MonthlyUsers:    5,
			}
			if !reflect.DeepEqual(gotStats, wantStats) {
				t.Errorf("UserStatistics() gotStats = \n%+v\nwant\n%+v", gotStats, wantStats)
			}
		})

		/*
			R30UsersV2 counts the number of 30 day retained users, defined as users that:
			- Appear more than once in the past 60 days
			- Have more than 30 days between the most and least recent appearances that occurred in the past 60 days.
			most recent -> neueste
			least recent -> älteste

		*/
		t.Run("R30UsersV2 tests", func(t *testing.T) {
			// generate some data
			for i := 100; i > 0; i-- {
				mustUpdateDeviceLastSeen(t, ctx, db, "user1", time.Now().AddDate(0, 0, -i))
				mustUpdateDeviceLastSeen(t, ctx, db, "user5", time.Now().AddDate(0, 0, -i))
				startTime := time.Now().AddDate(0, 0, -i)
				err := statsDB.UpdateUserDailyVisits(ctx, nil, startTime, startTime.Truncate(time.Hour*24).Add(time.Hour))
				if err != nil {
					t.Fatalf("unable to update daily visits stats: %v", err)
				}
			}
			gotStats, _, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			wantStats := &types.UserStatistics{
				RegisteredUsersByType: map[string]int64{
					"native":  3,
					"bridged": 1,
				},
				R30Users: map[string]int64{
					"all":      2,
					"android":  1,
					"electron": 1,
				},
				R30UsersV2: map[string]int64{
					"ios":      0,
					"android":  1,
					"web":      1,
					"electron": 0,
					"all":      2,
				},
				AllUsers:        6,
				NonBridgedUsers: 5,
				DailyUsers:      3,
				MonthlyUsers:    5,
			}
			if !reflect.DeepEqual(gotStats, wantStats) {
				t.Errorf("UserStatistics() gotStats = \n%+v\nwant\n%+v", gotStats, wantStats)
			}
		})
	})

}
<commit_msg>Fix flakey stats tests<commit_after>package tables_test

import (
	"context"
	"database/sql"
	"fmt"
	"reflect"
	"testing"
	"time"

	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"

	"github.com/matrix-org/dendrite/internal/sqlutil"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/dendrite/test"
	"github.com/matrix-org/dendrite/userapi/api"
	"github.com/matrix-org/dendrite/userapi/storage/postgres"
	"github.com/matrix-org/dendrite/userapi/storage/sqlite3"
	"github.com/matrix-org/dendrite/userapi/storage/tables"
	"github.com/matrix-org/dendrite/userapi/types"
)

func mustMakeDBs(t *testing.T, dbType test.DBType) (
	*sql.DB, tables.AccountsTable, tables.DevicesTable, tables.StatsTable, func(),
) {
	t.Helper()

	var (
		accTable   tables.AccountsTable
		devTable   tables.DevicesTable
		statsTable tables.StatsTable
		err        error
	)

	connStr, close := test.PrepareDBConnectionString(t, dbType)
	db, err := sqlutil.Open(&config.DatabaseOptions{
		ConnectionString: config.DataSource(connStr),
	}, nil)
	if err != nil {
		t.Fatalf("failed to open db: %s", err)
	}

	switch dbType {
	case test.DBTypeSQLite:
		accTable, err = sqlite3.NewSQLiteAccountsTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to create acc db: %v", err)
		}
		devTable, err = sqlite3.NewSQLiteDevicesTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to open device db: %v", err)
		}
		statsTable, err = sqlite3.NewSQLiteStatsTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to open stats db: %v", err)
		}
	case test.DBTypePostgres:
		accTable, err = postgres.NewPostgresAccountsTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to create acc db: %v", err)
		}
		devTable, err = postgres.NewPostgresDevicesTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to open device db: %v", err)
		}
		statsTable, err = postgres.NewPostgresStatsTable(db, "localhost")
		if err != nil {
			t.Fatalf("unable to open stats db: %v", err)
		}
	}

	return db, accTable, devTable, statsTable, close
}

func mustMakeAccountAndDevice(
	t *testing.T,
	ctx context.Context,
	accDB tables.AccountsTable,
	devDB tables.DevicesTable,
	localpart string,
	accType api.AccountType,
	userAgent string,
) {
	t.Helper()

	appServiceID := ""
	if accType == api.AccountTypeAppService {
		appServiceID = util.RandomString(16)
	}

	_, err := accDB.InsertAccount(ctx, nil, localpart, "", appServiceID, accType)
	if err != nil {
		t.Fatalf("unable to create account: %v", err)
	}
	_, err = devDB.InsertDevice(ctx, nil, "deviceID", localpart, util.RandomString(16), nil, "", userAgent)
	if err != nil {
		t.Fatalf("unable to create device: %v", err)
	}
}

func mustUpdateDeviceLastSeen(
	t *testing.T,
	ctx context.Context,
	db *sql.DB,
	localpart string,
	timestamp time.Time,
) {
	t.Helper()
	_, err := db.ExecContext(ctx, "UPDATE userapi_devices SET last_seen_ts = $1 WHERE localpart = $2", gomatrixserverlib.AsTimestamp(timestamp), localpart)
	if err != nil {
		t.Fatalf("unable to update device last seen")
	}
}

func mustUserUpdateRegistered(
	t *testing.T,
	ctx context.Context,
	db *sql.DB,
	localpart string,
	timestamp time.Time,
) {
	_, err := db.ExecContext(ctx, "UPDATE userapi_accounts SET created_ts = $1 WHERE localpart = $2", gomatrixserverlib.AsTimestamp(timestamp), localpart)
	if err != nil {
		t.Fatalf("unable to update device last seen")
	}
}

// These tests must run sequentially, as they build up on each other
func Test_UserStatistics(t *testing.T) {

	ctx := context.Background()
	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
		db, accDB, devDB, statsDB, close := mustMakeDBs(t, dbType)
		defer close()
		wantType := "SQLite"
		if dbType == test.DBTypePostgres {
			wantType = "Postgres"
		}

		t.Run(fmt.Sprintf("want %s database engine", wantType), func(t *testing.T) {
			_, gotDB, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			if wantType != gotDB.Engine { // can't use DeepEqual, as the Version might differ
				t.Errorf("UserStatistics() got DB engine = %+v, want %s", gotDB.Engine, wantType)
			}
		})

		t.Run("Want Users", func(t *testing.T) {
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user1", api.AccountTypeUser, "Element Android")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user2", api.AccountTypeUser, "Element iOS")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user3", api.AccountTypeUser, "Element web")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user4", api.AccountTypeGuest, "Element Electron")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user5", api.AccountTypeAdmin, "gecko")
			mustMakeAccountAndDevice(t, ctx, accDB, devDB, "user6", api.AccountTypeAppService, "gecko")
			gotStats, _, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			wantStats := &types.UserStatistics{
				RegisteredUsersByType: map[string]int64{
					"native":  4,
					"guest":   1,
					"bridged": 1,
				},
				R30Users: map[string]int64{},
				R30UsersV2: map[string]int64{
					"ios":      0,
					"android":  0,
					"web":      0,
					"electron": 0,
					"all":      0,
				},
				AllUsers:        6,
				NonBridgedUsers: 5,
				DailyUsers:      6,
				MonthlyUsers:    6,
			}
			if !reflect.DeepEqual(gotStats, wantStats) {
				t.Errorf("UserStatistics() gotStats = \n%+v\nwant\n%+v", gotStats, wantStats)
			}
		})

		t.Run("Users not active for one/two month", func(t *testing.T) {
			mustUpdateDeviceLastSeen(t, ctx, db, "user1", time.Now().AddDate(0, -2, 0))
			mustUpdateDeviceLastSeen(t, ctx, db, "user2", time.Now().AddDate(0, -1, 0))
			gotStats, _, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			wantStats := &types.UserStatistics{
				RegisteredUsersByType: map[string]int64{
					"native":  4,
					"guest":   1,
					"bridged": 1,
				},
				R30Users: map[string]int64{},
				R30UsersV2: map[string]int64{
					"ios":      0,
					"android":  0,
					"web":      0,
					"electron": 0,
					"all":      0,
				},
				AllUsers:        6,
				NonBridgedUsers: 5,
				DailyUsers:      4,
				MonthlyUsers:    4,
			}
			if !reflect.DeepEqual(gotStats, wantStats) {
				t.Errorf("UserStatistics() gotStats = \n%+v\nwant\n%+v", gotStats, wantStats)
			}
		})

		/* R30Users counts the number of 30 day retained users, defined as:
		- Users who have created their accounts more than 30 days ago
		- Where last seen at most 30 days ago
		- Where account creation and last_seen are > 30 days apart
		*/
		t.Run("R30Users tests", func(t *testing.T) {
			mustUserUpdateRegistered(t, ctx, db, "user1", time.Now().AddDate(0, -2, 0))
			mustUpdateDeviceLastSeen(t, ctx, db, "user1", time.Now())
			mustUserUpdateRegistered(t, ctx, db, "user4", time.Now().AddDate(0, -2, 0))
			mustUpdateDeviceLastSeen(t, ctx, db, "user4", time.Now())
			startTime := time.Now().AddDate(0, 0, -2)
			err := statsDB.UpdateUserDailyVisits(ctx, nil, startTime, startTime.Truncate(time.Hour*24))
			if err != nil {
				t.Fatalf("unable to update daily visits stats: %v", err)
			}

			gotStats, _, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			wantStats := &types.UserStatistics{
				RegisteredUsersByType: map[string]int64{
					"native":  3,
					"bridged": 1,
				},
				R30Users: map[string]int64{
					"all":      2,
					"android":  1,
					"electron": 1,
				},
				R30UsersV2: map[string]int64{
					"ios":      0,
					"android":  0,
					"web":      0,
					"electron": 0,
					"all":      0,
				},
				AllUsers:        6,
				NonBridgedUsers: 5,
				DailyUsers:      5,
				MonthlyUsers:    5,
			}
			if !reflect.DeepEqual(gotStats, wantStats) {
				t.Errorf("UserStatistics() gotStats = \n%+v\nwant\n%+v", gotStats, wantStats)
			}
		})

		/*
			R30UsersV2 counts the number of 30 day retained users, defined as users that:
			- Appear more than once in the past 60 days
			- Have more than 30 days between the most and least recent appearances that occurred in the past 60 days.
			most recent -> neueste
			least recent -> älteste

		*/
		t.Run("R30UsersV2 tests", func(t *testing.T) {
			// generate some data
			for i := 100; i > 0; i-- {
				mustUpdateDeviceLastSeen(t, ctx, db, "user1", time.Now().AddDate(0, 0, -i))
				mustUpdateDeviceLastSeen(t, ctx, db, "user5", time.Now().AddDate(0, 0, -i))
				startTime := time.Now().AddDate(0, 0, -i)
				err := statsDB.UpdateUserDailyVisits(ctx, nil, startTime, startTime.Truncate(time.Hour*24))
				if err != nil {
					t.Fatalf("unable to update daily visits stats: %v", err)
				}
			}
			gotStats, _, err := statsDB.UserStatistics(ctx, nil)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			wantStats := &types.UserStatistics{
				RegisteredUsersByType: map[string]int64{
					"native":  3,
					"bridged": 1,
				},
				R30Users: map[string]int64{
					"all":      2,
					"android":  1,
					"electron": 1,
				},
				R30UsersV2: map[string]int64{
					"ios":      0,
					"android":  1,
					"web":      1,
					"electron": 0,
					"all":      2,
				},
				AllUsers:        6,
				NonBridgedUsers: 5,
				DailyUsers:      3,
				MonthlyUsers:    5,
			}
			if !reflect.DeepEqual(gotStats, wantStats) {
				t.Errorf("UserStatistics() gotStats = \n%+v\nwant\n%+v", gotStats, wantStats)
			}
		})
	})

}
<|endoftext|>


<commit_before>package trakerr

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

type EventTraceBuilder struct {
}

func (tb *EventTraceBuilder) GetEventTraces(err interface{}, depth int, skip int) []InnerStackTrace {
	if err == nil {
		return nil
	}

	var traces = []InnerStackTrace{}

	return tb.AddStackTrace(traces, err, depth, skip+1)
}

func (tb *EventTraceBuilder) AddStackTrace(traces []InnerStackTrace, err interface{}, depth int, skip int) []InnerStackTrace {
	var innerTrace = InnerStackTrace{}

	innerTrace.TraceLines = tb.GetTraceLines(err, depth, skip+1)
	innerTrace.Message = fmt.Sprint(err)
	innerTrace.Type_ = fmt.Sprintf("%T", err)

	traces = append(traces, innerTrace)
	return traces
}

func (tb *EventTraceBuilder) GetTraceLines(err interface{}, depth int, skip int) []StackTraceLine {
	var traceLines = []StackTraceLine{}
	var goPath = os.Getenv("GOPATH")
	var goFilePath = tb.FileErrorHandler(filepath.Abs(goPath))
	for i := 0; i < depth; i++ {
		pc, file, line, ok := runtime.Caller(skip + 1 + i)
		if !ok {
			break
		}
		var localFilePath = tb.FileErrorHandler(filepath.Abs(file))

		var function = runtime.FuncForPC(pc)
		stLine := StackTraceLine{}
		stLine.File = strings.TrimPrefix(strings.ToLower(localFilePath), strings.ToLower(goFilePath))
		stLine.Line = int32(line)
		stLine.Function = function.Name()
		traceLines = append(traceLines, stLine)
	}

	return traceLines
}

//FileErrorHandler ...
func (tb *EventTraceBuilder) FileErrorHandler(str string, er error) string {
	if er != nil {
		panic(er)
	}

	return str
}
<commit_msg>Fix to runtime path truncation.<commit_after>package trakerr

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
)

type EventTraceBuilder struct {
}

func (tb *EventTraceBuilder) GetEventTraces(err interface{}, depth int, skip int) []InnerStackTrace {
	if err == nil {
		return nil
	}

	var traces = []InnerStackTrace{}

	return tb.AddStackTrace(traces, err, depth, skip+1)
}

func (tb *EventTraceBuilder) AddStackTrace(traces []InnerStackTrace, err interface{}, depth int, skip int) []InnerStackTrace {
	var innerTrace = InnerStackTrace{}

	innerTrace.TraceLines = tb.GetTraceLines(err, depth, skip+1)
	innerTrace.Message = fmt.Sprint(err)
	innerTrace.Type_ = fmt.Sprintf("%T", err)

	traces = append(traces, innerTrace)
	return traces
}

func (tb *EventTraceBuilder) GetTraceLines(err interface{}, depth int, skip int) []StackTraceLine {
	var traceLines = []StackTraceLine{}
	var goPath = tb.FileErrorHandler(filepath.Abs(os.Getenv("GOPATH")))
	var goRuntime = tb.FileErrorHandler(filepath.Abs(os.Getenv("GOROOT")))
	for i := 0; i < depth; i++ {
		pc, file, line, ok := runtime.Caller(skip + 1 + i)
		if !ok {
			break
		}
		var localFilePath = tb.FileErrorHandler(filepath.Abs(file))

		var function = runtime.FuncForPC(pc)
		stLine := StackTraceLine{}
		finalstring := strings.TrimPrefix(strings.ToLower(localFilePath), strings.ToLower(goPath))
		finalstring = strings.TrimPrefix(finalstring, strings.ToLower(goRuntime))
		stLine.File = finalstring
		stLine.Line = int32(line)
		stLine.Function = function.Name()
		traceLines = append(traceLines, stLine)
	}

	return traceLines
}

//FileErrorHandler ...
func (tb *EventTraceBuilder) FileErrorHandler(str string, er error) string {
	if er != nil {
		panic(er)
	}

	return str
}
<|endoftext|>


<commit_before>package ubiquity

// In this file, we mainly cover crypto ubiquity.
import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"crypto/x509"
	"math"
)

// HashUbiquity represents a score for how ubiquitous a given hash
// algorithm is; the higher the score, the more preferable the algorithm
// is.
type HashUbiquity int

// KeyAlgoUbiquity represents a score for how ubiquitous a given
// public-key algorithm is; the higher the score, the more preferable
// the algorithm is.
type KeyAlgoUbiquity int

// SHA1 is ubiquitous. SHA2 is not supported on some legacy platforms.
// We consider MD2/MD5 is harmful and thus assign them lowest ubiquity.
const (
	UnknownHashUbiquity HashUbiquity = 0
	SHA2Ubiquity        HashUbiquity = 70
	SHA1Ubiquity        HashUbiquity = 100
	MD5Ubiquity         HashUbiquity = 0
	MD2Ubiquity         HashUbiquity = 0
)

// RSA and DSA are considered ubiquitous. ECDSA256 and ECDSA384 should be
// supported by TLS 1.2 and have limited support from TLS 1.0 and
// 1.1, based on RFC6460, but ECDSA521 is less well-supported as
// a standard.
const (
	RSAUbiquity         KeyAlgoUbiquity = 100
	DSAUbiquity         KeyAlgoUbiquity = 100
	ECDSA256Ubiquity    KeyAlgoUbiquity = 70
	ECDSA384Ubiquity    KeyAlgoUbiquity = 70
	ECDSA521Ubiquity    KeyAlgoUbiquity = 30
	UnknownAlgoUbiquity KeyAlgoUbiquity = 0
)

// hashUbiquity computes the ubiquity of the hash algorithm in the
// signature algorithm of a cert.
// SHA1 > SHA2 > MD > Others
func hashUbiquity(cert *x509.Certificate) HashUbiquity {
	switch cert.SignatureAlgorithm {
	case x509.ECDSAWithSHA1, x509.DSAWithSHA1, x509.SHA1WithRSA:
		return SHA1Ubiquity
	case x509.ECDSAWithSHA256, x509.ECDSAWithSHA384, x509.ECDSAWithSHA512,
		x509.DSAWithSHA256, x509.SHA256WithRSA, x509.SHA384WithRSA,
		x509.SHA512WithRSA:
		return SHA2Ubiquity
	case x509.MD5WithRSA, x509.MD2WithRSA:
		return MD5Ubiquity
	default:
		return UnknownHashUbiquity
	}
}

// keyAlgoUbiquity compute the ubiquity of the cert's public key algorithm
// RSA, DSA>ECDSA>Unknown
func keyAlgoUbiquity(cert *x509.Certificate) KeyAlgoUbiquity {
	switch cert.PublicKeyAlgorithm {
	case x509.ECDSA:
		switch cert.PublicKey.(*ecdsa.PublicKey).Curve {
		case elliptic.P256():
			return ECDSA256Ubiquity
		case elliptic.P384():
			return ECDSA384Ubiquity
		case elliptic.P521():
			return ECDSA521Ubiquity
		default:
			return UnknownAlgoUbiquity
		}
	case x509.RSA:
		if cert.PublicKey.(*rsa.PublicKey).N.BitLen() >= 1024 {
			return RSAUbiquity
		}
		return UnknownAlgoUbiquity
	case x509.DSA:
		return DSAUbiquity
	default:
		return UnknownAlgoUbiquity
	}
}

// ChainHashUbiquity scores a chain based on the hash algorithms used
// by the certificates in the chain.
func ChainHashUbiquity(chain []*x509.Certificate) HashUbiquity {
	ret := math.MaxInt32
	for _, cert := range chain {
		uscore := int(hashUbiquity(cert))
		if ret > uscore {
			ret = uscore
		}
	}
	return HashUbiquity(ret)
}

// ChainKeyAlgoUbiquity scores a chain based on the public-key algorithms
// used by the certificates in the chain.
func ChainKeyAlgoUbiquity(chain []*x509.Certificate) KeyAlgoUbiquity {
	ret := math.MaxInt32
	for _, cert := range chain {
		uscore := int(keyAlgoUbiquity(cert))
		if ret > uscore {
			ret = uscore
		}
	}
	return KeyAlgoUbiquity(ret)
}

// CompareChainHashUbiquity returns a positive, zero, or negative value
// if the hash ubiquity of the first chain is greater, equal, or less
// than the second chain.
func CompareChainHashUbiquity(chain1, chain2 []*x509.Certificate) int {
	hu1 := ChainHashUbiquity(chain1)
	hu2 := ChainHashUbiquity(chain2)
	return int(hu1) - int(hu2)
}

// CompareChainKeyAlgoUbiquity returns a positive, zero, or negative value
// if the public-key ubiquity of the first chain is greater, equal,
// or less than the second chain.
func CompareChainKeyAlgoUbiquity(chain1, chain2 []*x509.Certificate) int {
	kau1 := ChainKeyAlgoUbiquity(chain1)
	kau2 := ChainKeyAlgoUbiquity(chain2)
	return int(kau1) - int(kau2)
}

// CompareExpiryUbiquity ranks two certificate chains based on the exiry dates of intermediates and roots.
// Certs expire later are ranked higher than ones expire earlier. The ranking between chains are determined by
// the first pair of intermediates, scanned from the root level,  that ar ranked differently.
func CompareExpiryUbiquity(chain1, chain2 []*x509.Certificate) int {
	for i := 0; ; i++ {
		if i >= len(chain1) || i >= len(chain2) {
			break
		}
		// Compare expiry dates, later is ranked higher.
		t1 := &chain1[len(chain1)-1-i].NotAfter
		t2 := &chain2[len(chain2)-1-i].NotAfter
		r := compareTime(t1, t2)
		// Return when we find rank difference.
		if r != 0 {
			return r
		}
	}
	return 0
}
<commit_msg>Add expiry validation to ubiquity check<commit_after>package ubiquity

// In this file, we mainly cover crypto ubiquity.
import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"crypto/x509"
	"github.com/cloudflare/cfssl/helpers"
	"math"
)

// HashUbiquity represents a score for how ubiquitous a given hash
// algorithm is; the higher the score, the more preferable the algorithm
// is.
type HashUbiquity int

// KeyAlgoUbiquity represents a score for how ubiquitous a given
// public-key algorithm is; the higher the score, the more preferable
// the algorithm is.
type KeyAlgoUbiquity int

// SHA1 is ubiquitous. SHA2 is not supported on some legacy platforms.
// We consider MD2/MD5 is harmful and thus assign them lowest ubiquity.
const (
	UnknownHashUbiquity HashUbiquity = 0
	SHA2Ubiquity        HashUbiquity = 70
	SHA1Ubiquity        HashUbiquity = 100
	MD5Ubiquity         HashUbiquity = 0
	MD2Ubiquity         HashUbiquity = 0
)

// RSA and DSA are considered ubiquitous. ECDSA256 and ECDSA384 should be
// supported by TLS 1.2 and have limited support from TLS 1.0 and
// 1.1, based on RFC6460, but ECDSA521 is less well-supported as
// a standard.
const (
	RSAUbiquity         KeyAlgoUbiquity = 100
	DSAUbiquity         KeyAlgoUbiquity = 100
	ECDSA256Ubiquity    KeyAlgoUbiquity = 70
	ECDSA384Ubiquity    KeyAlgoUbiquity = 70
	ECDSA521Ubiquity    KeyAlgoUbiquity = 30
	UnknownAlgoUbiquity KeyAlgoUbiquity = 0
)

// hashUbiquity computes the ubiquity of the hash algorithm in the
// signature algorithm of a cert.
// SHA1 > SHA2 > MD > Others
func hashUbiquity(cert *x509.Certificate) HashUbiquity {
	switch cert.SignatureAlgorithm {
	case x509.ECDSAWithSHA1, x509.DSAWithSHA1, x509.SHA1WithRSA:
		return SHA1Ubiquity
	case x509.ECDSAWithSHA256, x509.ECDSAWithSHA384, x509.ECDSAWithSHA512,
		x509.DSAWithSHA256, x509.SHA256WithRSA, x509.SHA384WithRSA,
		x509.SHA512WithRSA:
		return SHA2Ubiquity
	case x509.MD5WithRSA, x509.MD2WithRSA:
		return MD5Ubiquity
	default:
		return UnknownHashUbiquity
	}
}

// keyAlgoUbiquity compute the ubiquity of the cert's public key algorithm
// RSA, DSA>ECDSA>Unknown
func keyAlgoUbiquity(cert *x509.Certificate) KeyAlgoUbiquity {
	switch cert.PublicKeyAlgorithm {
	case x509.ECDSA:
		switch cert.PublicKey.(*ecdsa.PublicKey).Curve {
		case elliptic.P256():
			return ECDSA256Ubiquity
		case elliptic.P384():
			return ECDSA384Ubiquity
		case elliptic.P521():
			return ECDSA521Ubiquity
		default:
			return UnknownAlgoUbiquity
		}
	case x509.RSA:
		if cert.PublicKey.(*rsa.PublicKey).N.BitLen() >= 1024 {
			return RSAUbiquity
		}
		return UnknownAlgoUbiquity
	case x509.DSA:
		return DSAUbiquity
	default:
		return UnknownAlgoUbiquity
	}
}

// ChainHashUbiquity scores a chain based on the hash algorithms used
// by the certificates in the chain.
func ChainHashUbiquity(chain []*x509.Certificate) HashUbiquity {
	ret := math.MaxInt32
	for _, cert := range chain {
		uscore := int(hashUbiquity(cert))
		if ret > uscore {
			ret = uscore
		}
	}
	return HashUbiquity(ret)
}

// ChainKeyAlgoUbiquity scores a chain based on the public-key algorithms
// used by the certificates in the chain.
func ChainKeyAlgoUbiquity(chain []*x509.Certificate) KeyAlgoUbiquity {
	ret := math.MaxInt32
	for _, cert := range chain {
		uscore := int(keyAlgoUbiquity(cert))
		if ret > uscore {
			ret = uscore
		}
	}
	return KeyAlgoUbiquity(ret)
}

// CompareChainHashUbiquity returns a positive, zero, or negative value
// if the hash ubiquity of the first chain is greater, equal, or less
// than the second chain.
func CompareChainHashUbiquity(chain1, chain2 []*x509.Certificate) int {
	hu1 := ChainHashUbiquity(chain1)
	hu2 := ChainHashUbiquity(chain2)
	return int(hu1) - int(hu2)
}

// CompareChainKeyAlgoUbiquity returns a positive, zero, or negative value
// if the public-key ubiquity of the first chain is greater, equal,
// or less than the second chain.
func CompareChainKeyAlgoUbiquity(chain1, chain2 []*x509.Certificate) int {
	kau1 := ChainKeyAlgoUbiquity(chain1)
	kau2 := ChainKeyAlgoUbiquity(chain2)
	return int(kau1) - int(kau2)
}

// CompareExpiryUbiquity ranks two certificate chains based on the exiry dates of intermediates and roots.
// Certs expire later are ranked higher than ones expire earlier. The ranking between chains are determined by
// the first pair of intermediates, scanned from the root level,  that ar ranked differently.
func CompareExpiryUbiquity(chain1, chain2 []*x509.Certificate) int {
	for i := 0; ; i++ {
		if i >= len(chain1) || i >= len(chain2) {
			break
		}
		c1 := chain1[len(chain1)-1-i]
		c2 := chain2[len(chain2)-1-i]
		t1 := &c1.NotAfter
		t2 := &c2.NotAfter

		// Check if expiry dates valid. Return if one or other is invalid.
		// Otherwise rank by expiry date. Later is ranked higher.
		c1Valid := helpers.ValidExpiry(c1)
		c2Valid := helpers.ValidExpiry(c2)
		if c1Valid && !c2Valid {
			return 1
		}
		if !c1Valid && c2Valid {
			return -1
		}

		r := compareTime(t1, t2)
		// Return when we find rank difference.
		if r != 0 {
			return r
		}
	}
	return 0
}
<|endoftext|>