<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package azure

import (
	"encoding/base64"

	. "launchpad.net/gocheck"
	"launchpad.net/gwacl"

	"fmt"
	"launchpad.net/juju-core/instance"
	"net/http"
)

type InstanceSuite struct{}

var _ = Suite(new(InstanceSuite))

// makeHostedServiceDescriptor creates a HostedServiceDescriptor with the
// given service name.
func makeHostedServiceDescriptor(name string) *gwacl.HostedServiceDescriptor {
	labelBase64 := base64.StdEncoding.EncodeToString([]byte("label"))
	return &gwacl.HostedServiceDescriptor{ServiceName: name, Label: labelBase64}
}

func (*StorageSuite) TestId(c *C) {
	serviceName := "test-name"
	testService := makeHostedServiceDescriptor(serviceName)
	azInstance := azureInstance{*testService, nil}
	c.Check(azInstance.Id(), Equals, instance.Id(serviceName))
}

func (*StorageSuite) TestDNSName(c *C) {
	// An instance's DNS name is computed from its hosted-service name.
	host := "hostname"
	testService := makeHostedServiceDescriptor(host)
	azInstance := azureInstance{*testService, nil}
	dnsName, err := azInstance.DNSName()
	c.Assert(err, IsNil)
	c.Check(dnsName, Equals, host+"."+AZURE_DOMAIN_NAME)
}

func (*StorageSuite) TestWaitDNSName(c *C) {
	// An Azure instance gets its DNS name immediately, so there's no
	// waiting involved.
	host := "hostname"
	testService := makeHostedServiceDescriptor(host)
	azInstance := azureInstance{*testService, nil}
	dnsName, err := azInstance.WaitDNSName()
	c.Assert(err, IsNil)
	c.Check(dnsName, Equals, host+"."+AZURE_DOMAIN_NAME)
}

func makeRole(name string, endpoints ...gwacl.InputEndpoint) gwacl.Role {
	return gwacl.Role{
		RoleName: name,
		ConfigurationSets: []gwacl.ConfigurationSet{
			{
				ConfigurationSetType: gwacl.CONFIG_SET_NETWORK,
				InputEndpoints:       &endpoints,
			},
		},
	}
}

func makeDeployment(name string, roles ...gwacl.Role) gwacl.Deployment {
	return gwacl.Deployment{
		Name:     name,
		RoleList: roles,
	}
}

func makeInputEndpoint(port int, protocol string) gwacl.InputEndpoint {
	return gwacl.InputEndpoint{
		LocalPort: port,
		Name:      fmt.Sprintf("%s%d", protocol, port),
		Port:      port,
		Protocol:  protocol,
	}
}

func serialize(c *C, object gwacl.AzureObject) []byte {
	xml, err := object.Serialize()
	c.Assert(err, IsNil)
	return []byte(xml)
}

func preparePortChangeConversation(
	c *C, service *gwacl.HostedServiceDescriptor,
	deployments []gwacl.Deployment) []gwacl.DispatcherResponse {
	// Construct the series of responses to expected requests.
	responses := []gwacl.DispatcherResponse{
		// First, GetHostedServiceProperties
		gwacl.NewDispatcherResponse(
			serialize(c, &gwacl.HostedService{
				Deployments:             deployments,
				HostedServiceDescriptor: *service,
				XMLNS: gwacl.XMLNS,
			}),
			http.StatusOK, nil),
	}
	for _, deployment := range deployments {
		for _, role := range deployment.RoleList {
			// GetRole returns a PersistentVMRole.
			persistentRole := &gwacl.PersistentVMRole{
				XMLNS:             gwacl.XMLNS,
				RoleName:          role.RoleName,
				ConfigurationSets: role.ConfigurationSets,
			}
			responses = append(responses, gwacl.NewDispatcherResponse(
				serialize(c, persistentRole), http.StatusOK, nil))
			// UpdateRole expects a 200 response, that's all.
			responses = append(responses,
				gwacl.NewDispatcherResponse(nil, http.StatusOK, nil))
		}
	}
	return responses
}

type expectedRequest struct {
	method     string
	urlpattern string
}

func assertPortChangeConversation(c *C, record []*gwacl.X509Request, expected []expectedRequest) {
	c.Assert(record, HasLen, len(expected))
	for index, request := range record {
		c.Check(request.Method, Equals, expected[index].method)
		c.Check(request.URL, Matches, expected[index].urlpattern)
	}
}

func (*StorageSuite) TestOpenPorts(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	deployments := []gwacl.Deployment{
		makeDeployment("deployment-one", makeRole("role-one"), makeRole("role-two")),
		makeDeployment("deployment-two", makeRole("role-three")),
	}
	record := gwacl.PatchManagementAPIResponses(
		preparePortChangeConversation(c, service, deployments))
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.OpenPorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Assert(err, IsNil)
	assertPortChangeConversation(c, *record, []expectedRequest{
		{"GET", ".*/services/hostedservices/service-name[?].*"},   // GetHostedServiceProperties
		{"GET", ".*/deployments/deployment-one/roles/role-one"},   // GetRole
		{"PUT", ".*/deployments/deployment-one/roles/role-one"},   // UpdateRole
		{"GET", ".*/deployments/deployment-one/roles/role-two"},   // GetRole
		{"PUT", ".*/deployments/deployment-one/roles/role-two"},   // UpdateRole
		{"GET", ".*/deployments/deployment-two/roles/role-three"}, // GetRole
		{"PUT", ".*/deployments/deployment-two/roles/role-three"}, // UpdateRole
	})

	// A representative UpdateRole payload includes configuration for the
	// ports requested.
	role := &gwacl.PersistentVMRole{}
	err = role.Deserialize((*record)[2].Payload)
	c.Assert(err, IsNil)
	c.Check(
		*(role.ConfigurationSets[0].InputEndpoints),
		DeepEquals, []gwacl.InputEndpoint{
			makeInputEndpoint(79, "tcp"),
			makeInputEndpoint(587, "tcp"),
			makeInputEndpoint(9, "udp"),
		})
}

func (*StorageSuite) TestOpenPortsFailsWhenUnableToGetServiceProperties(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := []gwacl.DispatcherResponse{
		// GetHostedServiceProperties breaks.
		gwacl.NewDispatcherResponse(nil, http.StatusInternalServerError, nil),
	}
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.OpenPorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "GET request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 1)
}

func (*StorageSuite) TestOpenPortsFailsWhenUnableToGetRole(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	deployments := []gwacl.Deployment{
		makeDeployment("deployment-one", makeRole("role-one")),
	}
	responses := []gwacl.DispatcherResponse{
		// First, GetHostedServiceProperties
		gwacl.NewDispatcherResponse(
			serialize(c, &gwacl.HostedService{
				Deployments:             deployments,
				HostedServiceDescriptor: *service,
				XMLNS: gwacl.XMLNS,
			}),
			http.StatusOK, nil),
		// Second, GetRole fails
		gwacl.NewDispatcherResponse(
			nil, http.StatusInternalServerError, nil),
	}
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.OpenPorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "GET request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 2)
}

func (*StorageSuite) TestOpenPortsFailsWhenUnableToUpdateRole(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	deployments := []gwacl.Deployment{
		makeDeployment("deployment-one", makeRole("role-one")),
	}
	responses := []gwacl.DispatcherResponse{
		// First, GetHostedServiceProperties
		gwacl.NewDispatcherResponse(
			serialize(c, &gwacl.HostedService{
				Deployments:             deployments,
				HostedServiceDescriptor: *service,
				XMLNS: gwacl.XMLNS,
			}),
			http.StatusOK, nil),
		// Second, GetRole
		gwacl.NewDispatcherResponse(
			serialize(c, &gwacl.PersistentVMRole{
				XMLNS:    gwacl.XMLNS,
				RoleName: "role-one",
			}),
			http.StatusOK, nil),
		// Third, UpdateRole fails
		gwacl.NewDispatcherResponse(
			nil, http.StatusInternalServerError, nil),
	}
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.OpenPorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "PUT request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 3)
}

func (*StorageSuite) TestClosePorts(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	deployments := []gwacl.Deployment{
		makeDeployment("deployment-one",
			makeRole("role-one",
				makeInputEndpoint(587, "tcp"),
			),
			makeRole("role-two",
				makeInputEndpoint(79, "tcp"),
				makeInputEndpoint(9, "udp"),
			)),
		makeDeployment("deployment-two",
			makeRole("role-three",
				makeInputEndpoint(9, "tcp"),
				makeInputEndpoint(9, "udp"),
			)),
	}
	record := gwacl.PatchManagementAPIResponses(
		preparePortChangeConversation(c, service, deployments))
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.ClosePorts("machine-id", []instance.Port{{"tcp", 587}, {"udp", 9}})

	c.Assert(err, IsNil)
	assertPortChangeConversation(c, *record, []expectedRequest{
		{"GET", ".*/services/hostedservices/service-name[?].*"},   // GetHostedServiceProperties
		{"GET", ".*/deployments/deployment-one/roles/role-one"},   // GetRole
		{"PUT", ".*/deployments/deployment-one/roles/role-one"},   // UpdateRole
		{"GET", ".*/deployments/deployment-one/roles/role-two"},   // GetRole
		{"PUT", ".*/deployments/deployment-one/roles/role-two"},   // UpdateRole
		{"GET", ".*/deployments/deployment-two/roles/role-three"}, // GetRole
		{"PUT", ".*/deployments/deployment-two/roles/role-three"}, // UpdateRole
	})

	// The first UpdateRole removes all endpoints from the role's
	// configuration.
	roleOne := &gwacl.PersistentVMRole{}
	err = roleOne.Deserialize((*record)[2].Payload)
	c.Assert(err, IsNil)
	c.Check(roleOne.ConfigurationSets[0].InputEndpoints, IsNil)

	// The second UpdateRole removes all but 79/TCP.
	roleTwo := &gwacl.PersistentVMRole{}
	err = roleTwo.Deserialize((*record)[4].Payload)
	c.Assert(err, IsNil)
	c.Check(roleTwo.ConfigurationSets[0].InputEndpoints, DeepEquals,
		&[]gwacl.InputEndpoint{makeInputEndpoint(79, "tcp")})

	// The third UpdateRole removes all but 9/TCP.
	roleThree := &gwacl.PersistentVMRole{}
	err = roleThree.Deserialize((*record)[6].Payload)
	c.Assert(err, IsNil)
	c.Check(roleThree.ConfigurationSets[0].InputEndpoints, DeepEquals,
		&[]gwacl.InputEndpoint{makeInputEndpoint(9, "tcp")})
}

func (*StorageSuite) TestClosePortsFailsWhenUnableToGetServiceProperties(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := []gwacl.DispatcherResponse{
		// GetHostedServiceProperties breaks.
		gwacl.NewDispatcherResponse(nil, http.StatusInternalServerError, nil),
	}
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.ClosePorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "GET request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 1)
}

func (*StorageSuite) TestClosePortsFailsWhenUnableToGetRole(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	deployments := []gwacl.Deployment{
		makeDeployment("deployment-one", makeRole("role-one")),
	}
	responses := []gwacl.DispatcherResponse{
		// First, GetHostedServiceProperties
		gwacl.NewDispatcherResponse(
			serialize(c, &gwacl.HostedService{
				Deployments:             deployments,
				HostedServiceDescriptor: *service,
				XMLNS: gwacl.XMLNS,
			}),
			http.StatusOK, nil),
		// Second, GetRole fails
		gwacl.NewDispatcherResponse(
			nil, http.StatusInternalServerError, nil),
	}
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.ClosePorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "GET request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 2)
}

func (*StorageSuite) TestClosePortsFailsWhenUnableToUpdateRole(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	deployments := []gwacl.Deployment{
		makeDeployment("deployment-one", makeRole("role-one")),
	}
	responses := []gwacl.DispatcherResponse{
		// First, GetHostedServiceProperties
		gwacl.NewDispatcherResponse(
			serialize(c, &gwacl.HostedService{
				Deployments:             deployments,
				HostedServiceDescriptor: *service,
				XMLNS: gwacl.XMLNS,
			}),
			http.StatusOK, nil),
		// Second, GetRole
		gwacl.NewDispatcherResponse(
			serialize(c, &gwacl.PersistentVMRole{
				XMLNS:    gwacl.XMLNS,
				RoleName: "role-one",
			}),
			http.StatusOK, nil),
		// Third, UpdateRole fails
		gwacl.NewDispatcherResponse(
			nil, http.StatusInternalServerError, nil),
	}
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.ClosePorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "PUT request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 3)
}
<commit_msg>Extract a ton of set-up.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package azure

import (
	"encoding/base64"

	. "launchpad.net/gocheck"
	"launchpad.net/gwacl"

	"fmt"
	"launchpad.net/juju-core/instance"
	"net/http"
)

type InstanceSuite struct{}

var _ = Suite(new(InstanceSuite))

// makeHostedServiceDescriptor creates a HostedServiceDescriptor with the
// given service name.
func makeHostedServiceDescriptor(name string) *gwacl.HostedServiceDescriptor {
	labelBase64 := base64.StdEncoding.EncodeToString([]byte("label"))
	return &gwacl.HostedServiceDescriptor{ServiceName: name, Label: labelBase64}
}

func (*StorageSuite) TestId(c *C) {
	serviceName := "test-name"
	testService := makeHostedServiceDescriptor(serviceName)
	azInstance := azureInstance{*testService, nil}
	c.Check(azInstance.Id(), Equals, instance.Id(serviceName))
}

func (*StorageSuite) TestDNSName(c *C) {
	// An instance's DNS name is computed from its hosted-service name.
	host := "hostname"
	testService := makeHostedServiceDescriptor(host)
	azInstance := azureInstance{*testService, nil}
	dnsName, err := azInstance.DNSName()
	c.Assert(err, IsNil)
	c.Check(dnsName, Equals, host+"."+AZURE_DOMAIN_NAME)
}

func (*StorageSuite) TestWaitDNSName(c *C) {
	// An Azure instance gets its DNS name immediately, so there's no
	// waiting involved.
	host := "hostname"
	testService := makeHostedServiceDescriptor(host)
	azInstance := azureInstance{*testService, nil}
	dnsName, err := azInstance.WaitDNSName()
	c.Assert(err, IsNil)
	c.Check(dnsName, Equals, host+"."+AZURE_DOMAIN_NAME)
}

func makeRole(name string, endpoints ...gwacl.InputEndpoint) gwacl.Role {
	return gwacl.Role{
		RoleName: name,
		ConfigurationSets: []gwacl.ConfigurationSet{
			{
				ConfigurationSetType: gwacl.CONFIG_SET_NETWORK,
				InputEndpoints:       &endpoints,
			},
		},
	}
}

func makeDeployment(name string, roles ...gwacl.Role) gwacl.Deployment {
	return gwacl.Deployment{
		Name:     name,
		RoleList: roles,
	}
}

func makeInputEndpoint(port int, protocol string) gwacl.InputEndpoint {
	return gwacl.InputEndpoint{
		LocalPort: port,
		Name:      fmt.Sprintf("%s%d", protocol, port),
		Port:      port,
		Protocol:  protocol,
	}
}

func serialize(c *C, object gwacl.AzureObject) []byte {
	xml, err := object.Serialize()
	c.Assert(err, IsNil)
	return []byte(xml)
}

func preparePortChangeConversation(
	c *C, service *gwacl.HostedServiceDescriptor,
	deployments ...gwacl.Deployment) []gwacl.DispatcherResponse {
	// Construct the series of responses to expected requests.
	responses := []gwacl.DispatcherResponse{
		// First, GetHostedServiceProperties
		gwacl.NewDispatcherResponse(
			serialize(c, &gwacl.HostedService{
				Deployments:             deployments,
				HostedServiceDescriptor: *service,
				XMLNS: gwacl.XMLNS,
			}),
			http.StatusOK, nil),
	}
	for _, deployment := range deployments {
		for _, role := range deployment.RoleList {
			// GetRole returns a PersistentVMRole.
			persistentRole := &gwacl.PersistentVMRole{
				XMLNS:             gwacl.XMLNS,
				RoleName:          role.RoleName,
				ConfigurationSets: role.ConfigurationSets,
			}
			responses = append(responses, gwacl.NewDispatcherResponse(
				serialize(c, persistentRole), http.StatusOK, nil))
			// UpdateRole expects a 200 response, that's all.
			responses = append(responses,
				gwacl.NewDispatcherResponse(nil, http.StatusOK, nil))
		}
	}
	return responses
}

// point is 1-indexed; it represents which request should fail.
func failPortChangeConversationAt(point int, responses []gwacl.DispatcherResponse) {
	responses[point-1] = gwacl.NewDispatcherResponse(
		nil, http.StatusInternalServerError, nil)
}

type expectedRequest struct {
	method     string
	urlpattern string
}

func assertPortChangeConversation(c *C, record []*gwacl.X509Request, expected []expectedRequest) {
	c.Assert(record, HasLen, len(expected))
	for index, request := range record {
		c.Check(request.Method, Equals, expected[index].method)
		c.Check(request.URL, Matches, expected[index].urlpattern)
	}
}

func (*StorageSuite) TestOpenPorts(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := preparePortChangeConversation(c, service,
		makeDeployment("deployment-one",
			makeRole("role-one"), makeRole("role-two")),
		makeDeployment("deployment-two",
			makeRole("role-three")))
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.OpenPorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Assert(err, IsNil)
	assertPortChangeConversation(c, *record, []expectedRequest{
		{"GET", ".*/services/hostedservices/service-name[?].*"},   // GetHostedServiceProperties
		{"GET", ".*/deployments/deployment-one/roles/role-one"},   // GetRole
		{"PUT", ".*/deployments/deployment-one/roles/role-one"},   // UpdateRole
		{"GET", ".*/deployments/deployment-one/roles/role-two"},   // GetRole
		{"PUT", ".*/deployments/deployment-one/roles/role-two"},   // UpdateRole
		{"GET", ".*/deployments/deployment-two/roles/role-three"}, // GetRole
		{"PUT", ".*/deployments/deployment-two/roles/role-three"}, // UpdateRole
	})

	// A representative UpdateRole payload includes configuration for the
	// ports requested.
	role := &gwacl.PersistentVMRole{}
	err = role.Deserialize((*record)[2].Payload)
	c.Assert(err, IsNil)
	c.Check(
		*(role.ConfigurationSets[0].InputEndpoints),
		DeepEquals, []gwacl.InputEndpoint{
			makeInputEndpoint(79, "tcp"),
			makeInputEndpoint(587, "tcp"),
			makeInputEndpoint(9, "udp"),
		})
}

func (*StorageSuite) TestOpenPortsFailsWhenUnableToGetServiceProperties(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := []gwacl.DispatcherResponse{
		// GetHostedServiceProperties breaks.
		gwacl.NewDispatcherResponse(nil, http.StatusInternalServerError, nil),
	}
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.OpenPorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "GET request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 1)
}

func (*StorageSuite) TestOpenPortsFailsWhenUnableToGetRole(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := preparePortChangeConversation(c, service,
		makeDeployment("deployment-one", makeRole("role-one")))
	failPortChangeConversationAt(2, responses) // 2nd request, GetRole
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.OpenPorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "GET request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 2)
}

func (*StorageSuite) TestOpenPortsFailsWhenUnableToUpdateRole(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := preparePortChangeConversation(c, service,
		makeDeployment("deployment-one", makeRole("role-one")))
	failPortChangeConversationAt(3, responses) // 3rd request, UpdateRole
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.OpenPorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "PUT request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 3)
}

func (*StorageSuite) TestClosePorts(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := preparePortChangeConversation(c, service,
		makeDeployment("deployment-one",
			makeRole("role-one",
				makeInputEndpoint(587, "tcp"),
			),
			makeRole("role-two",
				makeInputEndpoint(79, "tcp"),
				makeInputEndpoint(9, "udp"),
			)),
		makeDeployment("deployment-two",
			makeRole("role-three",
				makeInputEndpoint(9, "tcp"),
				makeInputEndpoint(9, "udp"),
			)))
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.ClosePorts("machine-id",
		[]instance.Port{{"tcp", 587}, {"udp", 9}})

	c.Assert(err, IsNil)
	assertPortChangeConversation(c, *record, []expectedRequest{
		{"GET", ".*/services/hostedservices/service-name[?].*"},   // GetHostedServiceProperties
		{"GET", ".*/deployments/deployment-one/roles/role-one"},   // GetRole
		{"PUT", ".*/deployments/deployment-one/roles/role-one"},   // UpdateRole
		{"GET", ".*/deployments/deployment-one/roles/role-two"},   // GetRole
		{"PUT", ".*/deployments/deployment-one/roles/role-two"},   // UpdateRole
		{"GET", ".*/deployments/deployment-two/roles/role-three"}, // GetRole
		{"PUT", ".*/deployments/deployment-two/roles/role-three"}, // UpdateRole
	})

	// The first UpdateRole removes all endpoints from the role's
	// configuration.
	roleOne := &gwacl.PersistentVMRole{}
	err = roleOne.Deserialize((*record)[2].Payload)
	c.Assert(err, IsNil)
	c.Check(roleOne.ConfigurationSets[0].InputEndpoints, IsNil)

	// The second UpdateRole removes all but 79/TCP.
	roleTwo := &gwacl.PersistentVMRole{}
	err = roleTwo.Deserialize((*record)[4].Payload)
	c.Assert(err, IsNil)
	c.Check(roleTwo.ConfigurationSets[0].InputEndpoints, DeepEquals,
		&[]gwacl.InputEndpoint{makeInputEndpoint(79, "tcp")})

	// The third UpdateRole removes all but 9/TCP.
	roleThree := &gwacl.PersistentVMRole{}
	err = roleThree.Deserialize((*record)[6].Payload)
	c.Assert(err, IsNil)
	c.Check(roleThree.ConfigurationSets[0].InputEndpoints, DeepEquals,
		&[]gwacl.InputEndpoint{makeInputEndpoint(9, "tcp")})
}

func (*StorageSuite) TestClosePortsFailsWhenUnableToGetServiceProperties(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := []gwacl.DispatcherResponse{
		// GetHostedServiceProperties breaks.
		gwacl.NewDispatcherResponse(nil, http.StatusInternalServerError, nil),
	}
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.ClosePorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "GET request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 1)
}

func (*StorageSuite) TestClosePortsFailsWhenUnableToGetRole(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := preparePortChangeConversation(c, service,
		makeDeployment("deployment-one", makeRole("role-one")))
	failPortChangeConversationAt(2, responses) // 2nd request, GetRole
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.ClosePorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "GET request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 2)
}

func (*StorageSuite) TestClosePortsFailsWhenUnableToUpdateRole(c *C) {
	service := makeHostedServiceDescriptor("service-name")
	responses := preparePortChangeConversation(c, service,
		makeDeployment("deployment-one", makeRole("role-one")))
	failPortChangeConversationAt(3, responses) // 3rd request, UpdateRole
	record := gwacl.PatchManagementAPIResponses(responses)
	azInstance := azureInstance{*service, makeEnviron(c)}

	err := azInstance.ClosePorts("machine-id", []instance.Port{
		{"tcp", 79}, {"tcp", 587}, {"udp", 9},
	})

	c.Check(err, ErrorMatches, "PUT request failed [(]500: Internal Server Error[)]")
	c.Check(*record, HasLen, 3)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package etcd

import (
	"context"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apiserver/pkg/registry/generic"
	genericregistry "k8s.io/apiserver/pkg/registry/generic/registry"
	"k8s.io/apiserver/pkg/registry/rest"
	"k8s.io/kube-aggregator/pkg/apis/apiregistration"
	"k8s.io/kube-aggregator/pkg/registry/apiservice"
)

// rest implements a RESTStorage for API services against etcd
type REST struct {
	*genericregistry.Store
}

// NewREST returns a RESTStorage object that will work against API services.
func NewREST(scheme *runtime.Scheme, optsGetter generic.RESTOptionsGetter) *REST {
	strategy := apiservice.NewStrategy(scheme)
	store := &genericregistry.Store{
		NewFunc:                  func() runtime.Object { return &apiregistration.APIService{} },
		NewListFunc:              func() runtime.Object { return &apiregistration.APIServiceList{} },
		PredicateFunc:            apiservice.MatchAPIService,
		DefaultQualifiedResource: apiregistration.Resource("apiservices"),

		CreateStrategy: strategy,
		UpdateStrategy: strategy,
		DeleteStrategy: strategy,
	}
	options := &generic.StoreOptions{RESTOptions: optsGetter, AttrFunc: apiservice.GetAttrs}
	if err := store.CompleteWithOptions(options); err != nil {
		panic(err) // TODO: Propagate error up
	}
	return &REST{store}
}

// NewStatusREST makes a RESTStorage for status that has more limited options.
// It is based on the original REST so that we can share the same underlying store
func NewStatusREST(scheme *runtime.Scheme, rest *REST) *StatusREST {
	statusStore := *rest.Store
	statusStore.CreateStrategy = nil
	statusStore.DeleteStrategy = nil
	statusStore.UpdateStrategy = apiservice.NewStatusStrategy(scheme)
	return &StatusREST{store: &statusStore}
}

type StatusREST struct {
	store *genericregistry.Store
}

var _ = rest.Patcher(&StatusREST{})

func (r *StatusREST) New() runtime.Object {
	return &apiregistration.APIService{}
}

// Get retrieves the object from the storage. It is required to support Patch.
func (r *StatusREST) Get(ctx context.Context, name string, options *metav1.GetOptions) (runtime.Object, error) {
	return r.store.Get(ctx, name, options)
}

// Update alters the status subset of an object.
func (r *StatusREST) Update(ctx context.Context, name string, objInfo rest.UpdatedObjectInfo, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, forceAllowCreate bool) (runtime.Object, bool, error) {
	return r.store.Update(ctx, name, objInfo, createValidation, updateValidation, forceAllowCreate)
}
<commit_msg>Propagate forceAllowCreate as false to all subresources<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package etcd

import (
	"context"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apiserver/pkg/registry/generic"
	genericregistry "k8s.io/apiserver/pkg/registry/generic/registry"
	"k8s.io/apiserver/pkg/registry/rest"
	"k8s.io/kube-aggregator/pkg/apis/apiregistration"
	"k8s.io/kube-aggregator/pkg/registry/apiservice"
)

// rest implements a RESTStorage for API services against etcd
type REST struct {
	*genericregistry.Store
}

// NewREST returns a RESTStorage object that will work against API services.
func NewREST(scheme *runtime.Scheme, optsGetter generic.RESTOptionsGetter) *REST {
	strategy := apiservice.NewStrategy(scheme)
	store := &genericregistry.Store{
		NewFunc:                  func() runtime.Object { return &apiregistration.APIService{} },
		NewListFunc:              func() runtime.Object { return &apiregistration.APIServiceList{} },
		PredicateFunc:            apiservice.MatchAPIService,
		DefaultQualifiedResource: apiregistration.Resource("apiservices"),

		CreateStrategy: strategy,
		UpdateStrategy: strategy,
		DeleteStrategy: strategy,
	}
	options := &generic.StoreOptions{RESTOptions: optsGetter, AttrFunc: apiservice.GetAttrs}
	if err := store.CompleteWithOptions(options); err != nil {
		panic(err) // TODO: Propagate error up
	}
	return &REST{store}
}

// NewStatusREST makes a RESTStorage for status that has more limited options.
// It is based on the original REST so that we can share the same underlying store
func NewStatusREST(scheme *runtime.Scheme, rest *REST) *StatusREST {
	statusStore := *rest.Store
	statusStore.CreateStrategy = nil
	statusStore.DeleteStrategy = nil
	statusStore.UpdateStrategy = apiservice.NewStatusStrategy(scheme)
	return &StatusREST{store: &statusStore}
}

type StatusREST struct {
	store *genericregistry.Store
}

var _ = rest.Patcher(&StatusREST{})

func (r *StatusREST) New() runtime.Object {
	return &apiregistration.APIService{}
}

// Get retrieves the object from the storage. It is required to support Patch.
func (r *StatusREST) Get(ctx context.Context, name string, options *metav1.GetOptions) (runtime.Object, error) {
	return r.store.Get(ctx, name, options)
}

// Update alters the status subset of an object.
func (r *StatusREST) Update(ctx context.Context, name string, objInfo rest.UpdatedObjectInfo, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, forceAllowCreate bool) (runtime.Object, bool, error) {
	// We are explicitly setting forceAllowCreate to false in the call to the underlying storage because
	// subresources should never allow create on update.
	return r.store.Update(ctx, name, objInfo, createValidation, updateValidation, false)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Skaffold Authors

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

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

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

package schema

import (
	"fmt"
	"testing"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/util"
	"github.com/GoogleContainerTools/skaffold/testutil"
	"k8s.io/client-go/tools/clientcmd/api"
)

const (
	minimalConfig = ``

	simpleConfig = `
build:
  tagPolicy:
    gitCommit: {}
  artifacts:
  - image: example
deploy:
  kubectl: {}
`
	// This config has two tag policies set.
	invalidConfig = `
build:
  tagPolicy:
    sha256: {}
    gitCommit: {}
  artifacts:
  - image: example
deploy:
  name: example
`

	completeConfig = `
build:
  tagPolicy:
    sha256: {}
  artifacts:
  - image: image1
    context: ./examples/app1
    docker:
      dockerfile: Dockerfile.dev
  - image: image2
    context: ./examples/app2
    bazel:
      target: //:example.tar
  googleCloudBuild:
    projectId: ID
deploy:
  kubectl:
   manifests:
   - dep.yaml
   - svc.yaml
`
	minimalKanikoConfig = `
build:
  kaniko:
    buildContext:
      gcsBucket: demo
`
	completeKanikoConfig = `
build:
  kaniko:
    buildContext:
      gcsBucket: demo
    pullSecret: /secret.json
    pullSecretName: secret-name
    namespace: nskaniko
    timeout: 120m
`
	badConfig = "bad config"
)

func TestParseConfig(t *testing.T) {
	cleanup := testutil.SetupFakeKubernetesContext(t, api.Config{CurrentContext: "cluster1"})
	defer cleanup()

	var tests = []struct {
		apiVersion  string
		description string
		config      string
		expected    util.VersionedConfig
		badReader   bool
		shouldErr   bool
	}{
		{
			apiVersion:  latest.Version,
			description: "Minimal config",
			config:      minimalConfig,
			expected: config(
				withLocalBuild(
					withGitTagger(),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Simple config",
			config:      simpleConfig,
			expected: config(
				withLocalBuild(
					withGitTagger(),
					withDockerArtifact("example", ".", "Dockerfile"),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Complete config",
			config:      completeConfig,
			expected: config(
				withGoogleCloudBuild("ID",
					withShaTagger(),
					withDockerArtifact("image1", "./examples/app1", "Dockerfile.dev"),
					withBazelArtifact("image2", "./examples/app2", "//:example.tar"),
				),
				withKubectlDeploy("dep.yaml", "svc.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Minimal Kaniko config",
			config:      minimalKanikoConfig,
			expected: config(
				withKanikoBuild("demo", "kaniko-secret", "default", "", "20m",
					withGitTagger(),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Complete Kaniko config",
			config:      completeKanikoConfig,
			expected: config(
				withKanikoBuild("demo", "secret-name", "nskaniko", "/secret.json", "120m",
					withGitTagger(),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Bad config",
			config:      badConfig,
			shouldErr:   true,
		},
		{
			apiVersion:  latest.Version,
			description: "two taggers defined",
			config:      invalidConfig,
			shouldErr:   true,
		},
		{
			apiVersion:  "",
			description: "ApiVersion not specified",
			config:      minimalConfig,
			shouldErr:   true,
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			tmp, cleanup := testutil.NewTempDir(t)
			defer cleanup()

			yaml := fmt.Sprintf("apiVersion: %s\nkind: Config\n%s", test.apiVersion, test.config)
			tmp.Write("skaffold.yaml", yaml)

			cfg, err := ParseConfig(tmp.Path("skaffold.yaml"), true, false)

			testutil.CheckErrorAndDeepEqual(t, test.shouldErr, err, test.expected, cfg)
		})
	}
}

func config(ops ...func(*latest.SkaffoldPipeline)) *latest.SkaffoldPipeline {
	cfg := &latest.SkaffoldPipeline{APIVersion: latest.Version, Kind: "Config"}
	for _, op := range ops {
		op(cfg)
	}
	return cfg
}

func withLocalBuild(ops ...func(*latest.BuildConfig)) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		b := latest.BuildConfig{BuildType: latest.BuildType{LocalBuild: &latest.LocalBuild{}}}
		for _, op := range ops {
			op(&b)
		}
		cfg.Build = b
	}
}

func withGoogleCloudBuild(id string, ops ...func(*latest.BuildConfig)) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		b := latest.BuildConfig{BuildType: latest.BuildType{GoogleCloudBuild: &latest.GoogleCloudBuild{
			ProjectID:   id,
			DockerImage: "gcr.io/cloud-builders/docker",
		}}}
		for _, op := range ops {
			op(&b)
		}
		cfg.Build = b
	}
}

func withKanikoBuild(bucket, secretName, namespace, secret string, timeout string, ops ...func(*latest.BuildConfig)) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		b := latest.BuildConfig{BuildType: latest.BuildType{KanikoBuild: &latest.KanikoBuild{
			BuildContext: &latest.KanikoBuildContext{
				GCSBucket: bucket,
			},
			PullSecretName: secretName,
			Namespace:      namespace,
			PullSecret:     secret,
			Timeout:        timeout,
			Image:          constants.DefaultKanikoImage,
		}}}
		for _, op := range ops {
			op(&b)
		}
		cfg.Build = b
	}
}

func withKubectlDeploy(manifests ...string) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		cfg.Deploy = latest.DeployConfig{
			DeployType: latest.DeployType{
				KubectlDeploy: &latest.KubectlDeploy{
					Manifests: manifests,
				},
			},
		}
	}
}

func withHelmDeploy() func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		cfg.Deploy = latest.DeployConfig{
			DeployType: latest.DeployType{
				HelmDeploy: &latest.HelmDeploy{},
			},
		}
	}
}

func withDockerArtifact(image, workspace, dockerfile string) func(*latest.BuildConfig) {
	return func(cfg *latest.BuildConfig) {
		cfg.Artifacts = append(cfg.Artifacts, &latest.Artifact{
			ImageName: image,
			Workspace: workspace,
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					DockerfilePath: dockerfile,
				},
			},
		})
	}
}

func withBazelArtifact(image, workspace, target string) func(*latest.BuildConfig) {
	return func(cfg *latest.BuildConfig) {
		cfg.Artifacts = append(cfg.Artifacts, &latest.Artifact{
			ImageName: image,
			Workspace: workspace,
			ArtifactType: latest.ArtifactType{
				BazelArtifact: &latest.BazelArtifact{
					BuildTarget: target,
				},
			},
		})
	}
}

func withTagPolicy(tagPolicy latest.TagPolicy) func(*latest.BuildConfig) {
	return func(cfg *latest.BuildConfig) { cfg.TagPolicy = tagPolicy }
}

func withGitTagger() func(*latest.BuildConfig) {
	return withTagPolicy(latest.TagPolicy{GitTagger: &latest.GitTagger{}})
}

func withShaTagger() func(*latest.BuildConfig) {
	return withTagPolicy(latest.TagPolicy{ShaTagger: &latest.ShaTagger{}})
}

func withProfiles(profiles ...latest.Profile) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		cfg.Profiles = profiles
	}
}

func TestUpgradeToNextVersion(t *testing.T) {
	for i, schemaVersion := range schemaVersions[0 : len(schemaVersions)-2] {
		from := schemaVersion
		to := schemaVersions[i+1]
		description := fmt.Sprintf("Upgrade from %s to %s", from.apiVersion, to.apiVersion)

		t.Run(description, func(t *testing.T) {
			factory, _ := schemaVersions.Find(from.apiVersion)
			newer, err := factory().Upgrade()

			testutil.CheckErrorAndDeepEqual(t, false, err, to.apiVersion, newer.GetVersion())
		})
	}
}

func TestCantUpgradeFromLatestVersion(t *testing.T) {
	factory, present := schemaVersions.Find(latest.Version)
	testutil.CheckDeepEqual(t, true, present)

	_, err := factory().Upgrade()
	testutil.CheckError(t, true, err)
}
<commit_msg>Add a failing test that shows that Upgrade forgets to set default values<commit_after>/*
Copyright 2018 The Skaffold Authors

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

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

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

package schema

import (
	"fmt"
	"testing"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/util"
	"github.com/GoogleContainerTools/skaffold/testutil"
	"k8s.io/client-go/tools/clientcmd/api"
)

const (
	minimalConfig = ``

	simpleConfig = `
build:
  tagPolicy:
    gitCommit: {}
  artifacts:
  - image: example
deploy:
  kubectl: {}
`
	// This config has two tag policies set.
	invalidConfig = `
build:
  tagPolicy:
    sha256: {}
    gitCommit: {}
  artifacts:
  - image: example
deploy:
  name: example
`

	completeConfig = `
build:
  tagPolicy:
    sha256: {}
  artifacts:
  - image: image1
    context: ./examples/app1
    docker:
      dockerfile: Dockerfile.dev
  - image: image2
    context: ./examples/app2
    bazel:
      target: //:example.tar
  googleCloudBuild:
    projectId: ID
deploy:
  kubectl:
   manifests:
   - dep.yaml
   - svc.yaml
`
	minimalKanikoConfig = `
build:
  kaniko:
    buildContext:
      gcsBucket: demo
`
	completeKanikoConfig = `
build:
  kaniko:
    buildContext:
      gcsBucket: demo
    pullSecret: /secret.json
    pullSecretName: secret-name
    namespace: nskaniko
    timeout: 120m
`
	badConfig = "bad config"
)

func TestParseConfig(t *testing.T) {
	cleanup := testutil.SetupFakeKubernetesContext(t, api.Config{CurrentContext: "cluster1"})
	defer cleanup()

	var tests = []struct {
		apiVersion  string
		description string
		config      string
		expected    util.VersionedConfig
		badReader   bool
		shouldErr   bool
	}{
		{
			apiVersion:  latest.Version,
			description: "Minimal config",
			config:      minimalConfig,
			expected: config(
				withLocalBuild(
					withGitTagger(),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  "skaffold/v1alpha1",
			description: "Old minimal config",
			config:      minimalConfig,
			expected: config(
				withLocalBuild(
					withGitTagger(),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Simple config",
			config:      simpleConfig,
			expected: config(
				withLocalBuild(
					withGitTagger(),
					withDockerArtifact("example", ".", "Dockerfile"),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Complete config",
			config:      completeConfig,
			expected: config(
				withGoogleCloudBuild("ID",
					withShaTagger(),
					withDockerArtifact("image1", "./examples/app1", "Dockerfile.dev"),
					withBazelArtifact("image2", "./examples/app2", "//:example.tar"),
				),
				withKubectlDeploy("dep.yaml", "svc.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Minimal Kaniko config",
			config:      minimalKanikoConfig,
			expected: config(
				withKanikoBuild("demo", "kaniko-secret", "default", "", "20m",
					withGitTagger(),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Complete Kaniko config",
			config:      completeKanikoConfig,
			expected: config(
				withKanikoBuild("demo", "secret-name", "nskaniko", "/secret.json", "120m",
					withGitTagger(),
				),
				withKubectlDeploy("k8s/*.yaml"),
			),
		},
		{
			apiVersion:  latest.Version,
			description: "Bad config",
			config:      badConfig,
			shouldErr:   true,
		},
		{
			apiVersion:  latest.Version,
			description: "two taggers defined",
			config:      invalidConfig,
			shouldErr:   true,
		},
		{
			apiVersion:  "",
			description: "ApiVersion not specified",
			config:      minimalConfig,
			shouldErr:   true,
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			tmp, cleanup := testutil.NewTempDir(t)
			defer cleanup()

			yaml := fmt.Sprintf("apiVersion: %s\nkind: Config\n%s", test.apiVersion, test.config)
			tmp.Write("skaffold.yaml", yaml)

			cfg, err := ParseConfig(tmp.Path("skaffold.yaml"), true, true)

			testutil.CheckErrorAndDeepEqual(t, test.shouldErr, err, test.expected, cfg)
		})
	}
}

func config(ops ...func(*latest.SkaffoldPipeline)) *latest.SkaffoldPipeline {
	cfg := &latest.SkaffoldPipeline{APIVersion: latest.Version, Kind: "Config"}
	for _, op := range ops {
		op(cfg)
	}
	return cfg
}

func withLocalBuild(ops ...func(*latest.BuildConfig)) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		b := latest.BuildConfig{BuildType: latest.BuildType{LocalBuild: &latest.LocalBuild{}}}
		for _, op := range ops {
			op(&b)
		}
		cfg.Build = b
	}
}

func withGoogleCloudBuild(id string, ops ...func(*latest.BuildConfig)) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		b := latest.BuildConfig{BuildType: latest.BuildType{GoogleCloudBuild: &latest.GoogleCloudBuild{
			ProjectID:   id,
			DockerImage: "gcr.io/cloud-builders/docker",
		}}}
		for _, op := range ops {
			op(&b)
		}
		cfg.Build = b
	}
}

func withKanikoBuild(bucket, secretName, namespace, secret string, timeout string, ops ...func(*latest.BuildConfig)) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		b := latest.BuildConfig{BuildType: latest.BuildType{KanikoBuild: &latest.KanikoBuild{
			BuildContext: &latest.KanikoBuildContext{
				GCSBucket: bucket,
			},
			PullSecretName: secretName,
			Namespace:      namespace,
			PullSecret:     secret,
			Timeout:        timeout,
			Image:          constants.DefaultKanikoImage,
		}}}
		for _, op := range ops {
			op(&b)
		}
		cfg.Build = b
	}
}

func withKubectlDeploy(manifests ...string) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		cfg.Deploy = latest.DeployConfig{
			DeployType: latest.DeployType{
				KubectlDeploy: &latest.KubectlDeploy{
					Manifests: manifests,
				},
			},
		}
	}
}

func withHelmDeploy() func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		cfg.Deploy = latest.DeployConfig{
			DeployType: latest.DeployType{
				HelmDeploy: &latest.HelmDeploy{},
			},
		}
	}
}

func withDockerArtifact(image, workspace, dockerfile string) func(*latest.BuildConfig) {
	return func(cfg *latest.BuildConfig) {
		cfg.Artifacts = append(cfg.Artifacts, &latest.Artifact{
			ImageName: image,
			Workspace: workspace,
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					DockerfilePath: dockerfile,
				},
			},
		})
	}
}

func withBazelArtifact(image, workspace, target string) func(*latest.BuildConfig) {
	return func(cfg *latest.BuildConfig) {
		cfg.Artifacts = append(cfg.Artifacts, &latest.Artifact{
			ImageName: image,
			Workspace: workspace,
			ArtifactType: latest.ArtifactType{
				BazelArtifact: &latest.BazelArtifact{
					BuildTarget: target,
				},
			},
		})
	}
}

func withTagPolicy(tagPolicy latest.TagPolicy) func(*latest.BuildConfig) {
	return func(cfg *latest.BuildConfig) { cfg.TagPolicy = tagPolicy }
}

func withGitTagger() func(*latest.BuildConfig) {
	return withTagPolicy(latest.TagPolicy{GitTagger: &latest.GitTagger{}})
}

func withShaTagger() func(*latest.BuildConfig) {
	return withTagPolicy(latest.TagPolicy{ShaTagger: &latest.ShaTagger{}})
}

func withProfiles(profiles ...latest.Profile) func(*latest.SkaffoldPipeline) {
	return func(cfg *latest.SkaffoldPipeline) {
		cfg.Profiles = profiles
	}
}

func TestUpgradeToNextVersion(t *testing.T) {
	for i, schemaVersion := range schemaVersions[0 : len(schemaVersions)-2] {
		from := schemaVersion
		to := schemaVersions[i+1]
		description := fmt.Sprintf("Upgrade from %s to %s", from.apiVersion, to.apiVersion)

		t.Run(description, func(t *testing.T) {
			factory, _ := schemaVersions.Find(from.apiVersion)
			newer, err := factory().Upgrade()

			testutil.CheckErrorAndDeepEqual(t, false, err, to.apiVersion, newer.GetVersion())
		})
	}
}

func TestCantUpgradeFromLatestVersion(t *testing.T) {
	factory, present := schemaVersions.Find(latest.Version)
	testutil.CheckDeepEqual(t, true, present)

	_, err := factory().Upgrade()
	testutil.CheckError(t, true, err)
}
<|endoftext|>


<commit_before>package cache

import (
	"context"
	"encoding/hex"
	"flag"
	"hash/fnv"
	"sync"
	"time"

	"github.com/bradfitz/gomemcache/memcache"
	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	instr "github.com/weaveworks/common/instrument"

	"github.com/grafana/loki/pkg/logqlmodel/stats"
	util_log "github.com/grafana/loki/pkg/util/log"
	"github.com/grafana/loki/pkg/util/math"
)

// MemcachedConfig is config to make a Memcached
type MemcachedConfig struct {
	Expiration time.Duration `yaml:"expiration"`

	BatchSize   int `yaml:"batch_size"`
	Parallelism int `yaml:"parallelism"`
}

// RegisterFlagsWithPrefix adds the flags required to config this to the given FlagSet
func (cfg *MemcachedConfig) RegisterFlagsWithPrefix(prefix, description string, f *flag.FlagSet) {
	f.DurationVar(&cfg.Expiration, prefix+"memcached.expiration", 0, description+"How long keys stay in the memcache.")
	f.IntVar(&cfg.BatchSize, prefix+"memcached.batchsize", 1024, description+"How many keys to fetch in each batch.")
	f.IntVar(&cfg.Parallelism, prefix+"memcached.parallelism", 100, description+"Maximum active requests to memcache.")
}

// Memcached type caches chunks in memcached
type Memcached struct {
	cfg       MemcachedConfig
	memcache  MemcachedClient
	name      string
	cacheType stats.CacheType

	requestDuration *instr.HistogramCollector

	wg      sync.WaitGroup
	inputCh chan *work

	logger log.Logger
}

// NewMemcached makes a new Memcached.
func NewMemcached(cfg MemcachedConfig, client MemcachedClient, name string, reg prometheus.Registerer, logger log.Logger, cacheType stats.CacheType) *Memcached {
	c := &Memcached{
		cfg:       cfg,
		memcache:  client,
		name:      name,
		logger:    logger,
		cacheType: cacheType,
		requestDuration: instr.NewHistogramCollector(
			promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
				Namespace: "loki",
				Name:      "memcache_request_duration_seconds",
				Help:      "Total time spent in seconds doing memcache requests.",
				// Memcached requests are very quick: smallest bucket is 16us, biggest is 1s
				Buckets:     prometheus.ExponentialBuckets(0.000016, 4, 8),
				ConstLabels: prometheus.Labels{"name": name},
			}, []string{"method", "status_code"}),
		),
	}

	if cfg.BatchSize == 0 || cfg.Parallelism == 0 {
		return c
	}

	c.inputCh = make(chan *work)
	c.wg.Add(cfg.Parallelism)

	for i := 0; i < cfg.Parallelism; i++ {
		go func() {
			for input := range c.inputCh {
				res := &result{
					batchID: input.batchID,
				}
				res.found, res.bufs, res.missed, res.err = c.fetch(input.ctx, input.keys)
				input.resultCh <- res
			}

			c.wg.Done()
		}()
	}

	return c
}

type work struct {
	keys     []string
	ctx      context.Context
	resultCh chan<- *result
	batchID  int // For ordering results.
}

type result struct {
	found   []string
	bufs    [][]byte
	missed  []string
	err     error
	batchID int // For ordering results.
}

func memcacheStatusCode(err error) string {
	// See https://godoc.org/github.com/bradfitz/gomemcache/memcache#pkg-variables
	switch err {
	case nil:
		return "200"
	case memcache.ErrCacheMiss:
		return "404"
	case memcache.ErrMalformedKey:
		return "400"
	default:
		return "500"
	}
}

// Fetch gets keys from the cache. The keys that are found must be in the order of the keys requested.
func (c *Memcached) Fetch(ctx context.Context, keys []string) (found []string, bufs [][]byte, missed []string, err error) {
	if c.cfg.BatchSize == 0 {
		found, bufs, missed, err = c.fetch(ctx, keys)
		return
	}

	start := time.Now()
	found, bufs, missed, err = c.fetchKeysBatched(ctx, keys)
	c.requestDuration.After(ctx, "Memcache.GetBatched", memcacheStatusCode(err), start)
	return
}

func (c *Memcached) fetch(ctx context.Context, keys []string) (found []string, bufs [][]byte, missed []string, err error) {
	var (
		start = time.Now()
		items map[string]*memcache.Item
	)
	items, err = c.memcache.GetMulti(keys)
	c.requestDuration.After(ctx, "Memcache.GetMulti", memcacheStatusCode(err), start)
	if err != nil {
		level.Error(util_log.WithContext(ctx, c.logger)).Log(
			"msg", "Failed to get keys from memcached",
			"keys requested", len(keys),
			"err", err,
		)
		return found, bufs, keys, err
	}

	for _, key := range keys {
		item, ok := items[key]
		if ok {
			found = append(found, key)
			bufs = append(bufs, item.Value)
		} else {
			missed = append(missed, key)
		}
	}
	return
}

func (c *Memcached) fetchKeysBatched(ctx context.Context, keys []string) (found []string, bufs [][]byte, missed []string, err error) {
	resultsCh := make(chan *result)
	batchSize := c.cfg.BatchSize

	go func() {
		for i, j := 0, 0; i < len(keys); i += batchSize {
			batchKeys := keys[i:math.Min(i+batchSize, len(keys))]
			c.inputCh <- &work{
				keys:     batchKeys,
				ctx:      ctx,
				resultCh: resultsCh,
				batchID:  j,
			}
			j++
		}
	}()

	// Read all values from this channel to avoid blocking upstream.
	numResults := len(keys) / batchSize
	if len(keys)%batchSize != 0 {
		numResults++
	}

	// We need to order found by the input keys order.
	results := make([]*result, numResults)
	for i := 0; i < numResults; i++ {
		result := <-resultsCh
		results[result.batchID] = result
	}
	close(resultsCh)

	for _, result := range results {
		found = append(found, result.found...)
		bufs = append(bufs, result.bufs...)
		missed = append(missed, result.missed...)
		if result.err != nil {
			err = result.err
		}
	}

	return
}

// Store stores the key in the cache.
func (c *Memcached) Store(ctx context.Context, keys []string, bufs [][]byte) error {
	var err error
	for i := range keys {
		cacheErr := instr.CollectedRequest(ctx, "Memcache.Put", c.requestDuration, memcacheStatusCode, func(_ context.Context) error {
			item := memcache.Item{
				Key:        keys[i],
				Value:      bufs[i],
				Expiration: int32(c.cfg.Expiration.Seconds()),
			}
			return c.memcache.Set(&item)
		})
		if cacheErr != nil {
			level.Warn(c.logger).Log("msg", "failed to put to memcached", "name", c.name, "err", err)
			err = cacheErr
		}
	}
	return err
}

// Stop does nothing.
func (c *Memcached) Stop() {
	if c.inputCh == nil {
		return
	}

	close(c.inputCh)
	c.wg.Wait()
}

func (c *Memcached) GetCacheType() stats.CacheType {
	return c.cacheType
}

// HashKey hashes key into something you can store in memcached.
func HashKey(key string) string {
	hasher := fnv.New64a()
	_, _ = hasher.Write([]byte(key)) // This'll never error.

	// Hex because memcache errors for the bytes produced by the hash.
	return hex.EncodeToString(hasher.Sum(nil))
}
<commit_msg>[Enhancement] Removes memcached failure log (#6720)<commit_after>package cache

import (
	"context"
	"encoding/hex"
	"flag"
	"hash/fnv"
	"sync"
	"time"

	"github.com/bradfitz/gomemcache/memcache"
	"github.com/go-kit/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	instr "github.com/weaveworks/common/instrument"

	"github.com/grafana/loki/pkg/logqlmodel/stats"
	"github.com/grafana/loki/pkg/util/math"
)

// MemcachedConfig is config to make a Memcached
type MemcachedConfig struct {
	Expiration time.Duration `yaml:"expiration"`

	BatchSize   int `yaml:"batch_size"`
	Parallelism int `yaml:"parallelism"`
}

// RegisterFlagsWithPrefix adds the flags required to config this to the given FlagSet
func (cfg *MemcachedConfig) RegisterFlagsWithPrefix(prefix, description string, f *flag.FlagSet) {
	f.DurationVar(&cfg.Expiration, prefix+"memcached.expiration", 0, description+"How long keys stay in the memcache.")
	f.IntVar(&cfg.BatchSize, prefix+"memcached.batchsize", 1024, description+"How many keys to fetch in each batch.")
	f.IntVar(&cfg.Parallelism, prefix+"memcached.parallelism", 100, description+"Maximum active requests to memcache.")
}

// Memcached type caches chunks in memcached
type Memcached struct {
	cfg       MemcachedConfig
	memcache  MemcachedClient
	name      string
	cacheType stats.CacheType

	requestDuration *instr.HistogramCollector

	wg      sync.WaitGroup
	inputCh chan *work

	logger log.Logger
}

// NewMemcached makes a new Memcached.
func NewMemcached(cfg MemcachedConfig, client MemcachedClient, name string, reg prometheus.Registerer, logger log.Logger, cacheType stats.CacheType) *Memcached {
	c := &Memcached{
		cfg:       cfg,
		memcache:  client,
		name:      name,
		logger:    logger,
		cacheType: cacheType,
		requestDuration: instr.NewHistogramCollector(
			promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
				Namespace: "loki",
				Name:      "memcache_request_duration_seconds",
				Help:      "Total time spent in seconds doing memcache requests.",
				// Memcached requests are very quick: smallest bucket is 16us, biggest is 1s
				Buckets:     prometheus.ExponentialBuckets(0.000016, 4, 8),
				ConstLabels: prometheus.Labels{"name": name},
			}, []string{"method", "status_code"}),
		),
	}

	if cfg.BatchSize == 0 || cfg.Parallelism == 0 {
		return c
	}

	c.inputCh = make(chan *work)
	c.wg.Add(cfg.Parallelism)

	for i := 0; i < cfg.Parallelism; i++ {
		go func() {
			for input := range c.inputCh {
				res := &result{
					batchID: input.batchID,
				}
				res.found, res.bufs, res.missed, res.err = c.fetch(input.ctx, input.keys)
				input.resultCh <- res
			}

			c.wg.Done()
		}()
	}

	return c
}

type work struct {
	keys     []string
	ctx      context.Context
	resultCh chan<- *result
	batchID  int // For ordering results.
}

type result struct {
	found   []string
	bufs    [][]byte
	missed  []string
	err     error
	batchID int // For ordering results.
}

func memcacheStatusCode(err error) string {
	// See https://godoc.org/github.com/bradfitz/gomemcache/memcache#pkg-variables
	switch err {
	case nil:
		return "200"
	case memcache.ErrCacheMiss:
		return "404"
	case memcache.ErrMalformedKey:
		return "400"
	default:
		return "500"
	}
}

// Fetch gets keys from the cache. The keys that are found must be in the order of the keys requested.
func (c *Memcached) Fetch(ctx context.Context, keys []string) (found []string, bufs [][]byte, missed []string, err error) {
	if c.cfg.BatchSize == 0 {
		found, bufs, missed, err = c.fetch(ctx, keys)
		return
	}

	start := time.Now()
	found, bufs, missed, err = c.fetchKeysBatched(ctx, keys)
	c.requestDuration.After(ctx, "Memcache.GetBatched", memcacheStatusCode(err), start)
	return
}

func (c *Memcached) fetch(ctx context.Context, keys []string) (found []string, bufs [][]byte, missed []string, err error) {
	var (
		start = time.Now()
		items map[string]*memcache.Item
	)
	items, err = c.memcache.GetMulti(keys)
	c.requestDuration.After(ctx, "Memcache.GetMulti", memcacheStatusCode(err), start)
	if err != nil {
		return found, bufs, keys, err
	}

	for _, key := range keys {
		item, ok := items[key]
		if ok {
			found = append(found, key)
			bufs = append(bufs, item.Value)
		} else {
			missed = append(missed, key)
		}
	}
	return
}

func (c *Memcached) fetchKeysBatched(ctx context.Context, keys []string) (found []string, bufs [][]byte, missed []string, err error) {
	resultsCh := make(chan *result)
	batchSize := c.cfg.BatchSize

	go func() {
		for i, j := 0, 0; i < len(keys); i += batchSize {
			batchKeys := keys[i:math.Min(i+batchSize, len(keys))]
			c.inputCh <- &work{
				keys:     batchKeys,
				ctx:      ctx,
				resultCh: resultsCh,
				batchID:  j,
			}
			j++
		}
	}()

	// Read all values from this channel to avoid blocking upstream.
	numResults := len(keys) / batchSize
	if len(keys)%batchSize != 0 {
		numResults++
	}

	// We need to order found by the input keys order.
	results := make([]*result, numResults)
	for i := 0; i < numResults; i++ {
		result := <-resultsCh
		results[result.batchID] = result
	}
	close(resultsCh)

	for _, result := range results {
		found = append(found, result.found...)
		bufs = append(bufs, result.bufs...)
		missed = append(missed, result.missed...)
		if result.err != nil {
			err = result.err
		}
	}

	return
}

// Store stores the key in the cache.
func (c *Memcached) Store(ctx context.Context, keys []string, bufs [][]byte) error {
	var err error
	for i := range keys {
		cacheErr := instr.CollectedRequest(ctx, "Memcache.Put", c.requestDuration, memcacheStatusCode, func(_ context.Context) error {
			item := memcache.Item{
				Key:        keys[i],
				Value:      bufs[i],
				Expiration: int32(c.cfg.Expiration.Seconds()),
			}
			return c.memcache.Set(&item)
		})
		if cacheErr != nil {
			err = cacheErr
		}
	}
	return err
}

// Stop does nothing.
func (c *Memcached) Stop() {
	if c.inputCh == nil {
		return
	}

	close(c.inputCh)
	c.wg.Wait()
}

func (c *Memcached) GetCacheType() stats.CacheType {
	return c.cacheType
}

// HashKey hashes key into something you can store in memcached.
func HashKey(key string) string {
	hasher := fnv.New64a()
	_, _ = hasher.Write([]byte(key)) // This'll never error.

	// Hex because memcache errors for the bytes produced by the hash.
	return hex.EncodeToString(hasher.Sum(nil))
}
<|endoftext|>


<commit_before>package proxy

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"

	"github.com/mailgun/vulcand/engine"

	"github.com/mailgun/vulcand/Godeps/_workspace/src/github.com/mailgun/log"
	"github.com/mailgun/vulcand/Godeps/_workspace/src/github.com/mailgun/manners"
)

// srv contains all that is necessary to run the HTTP(s) server. server does not work on its own,
// it heavily depends on MuxServer and acts as its internal data structure.
type srv struct {
	defaultHost string
	mux         *mux
	srv         *manners.GracefulServer
	proxy       http.Handler
	listener    engine.Listener
	keyPairs    map[engine.HostKey]engine.KeyPair
	options     Options
	state       int
}

func (s *srv) GetFile() (*FileDescriptor, error) {
	if !s.hasListeners() || s.srv == nil {
		return nil, nil
	}
	file, err := s.srv.GetFile()
	if err != nil {
		return nil, err
	}
	return &FileDescriptor{
		File:    file,
		Address: s.listener.Address,
	}, nil
}

func (s *srv) String() string {
	return fmt.Sprintf("%s->srv(%v, %v)", s.mux, s.state, &s.listener)
}

func newSrv(m *mux, l engine.Listener) (*srv, error) {
	defaultHost := ""
	keyPairs := make(map[engine.HostKey]engine.KeyPair)
	for hk, h := range m.hosts {
		if h.Settings.KeyPair != nil {
			keyPairs[hk] = *h.Settings.KeyPair
		}
		if h.Settings.Default {
			defaultHost = hk.Name
		}
	}

	return &srv{
		mux:         m,
		proxy:       m.router,
		listener:    l,
		defaultHost: defaultHost,
		keyPairs:    keyPairs,
		state:       srvStateInit,
	}, nil
}

func (s *srv) deleteKeyPair(hk engine.HostKey) error {
	delete(s.keyPairs, hk)
	return s.reload()
}

func (s *srv) isTLS() bool {
	return s.listener.Protocol == engine.HTTPS
}

func (s *srv) upsertKeyPair(hk engine.HostKey, keyPair *engine.KeyPair) error {
	old, exists := s.keyPairs[hk]
	if exists && old.Equals(keyPair) {
		return nil
	}
	s.keyPairs[hk] = *keyPair
	return s.reload()
}

func (s *srv) setDefaultHost(host engine.Host) error {
	oldDefault := s.defaultHost
	if host.Settings.Default {
		s.defaultHost = host.Name
	}
	if oldDefault != s.defaultHost && s.isTLS() {
		return s.reload()
	}
	return nil
}

func (s *srv) isServing() bool {
	return s.state == srvStateActive
}

func (s *srv) hasListeners() bool {
	return s.state == srvStateActive || s.state == srvStateHijacked
}

func (s *srv) takeFile(f *FileDescriptor) error {
	log.Infof("%s takeFile %v", s, f)

	listener, err := f.ToListener()
	if err != nil {
		return err
	}

	if s.isTLS() {
		tcpListener, ok := listener.(*net.TCPListener)
		if !ok {
			return fmt.Errorf(`%s failed to take file descriptor - it is running in TLS mode so I need a TCP listener,
but the file descriptor that was given corresponded to a listener of type %T. More about file descriptor: %s`, listener, s, f)
		}
		config, err := newTLSConfig(s.keyPairs, s.defaultHost)
		if err != nil {
			return err
		}
		listener = manners.NewTLSListener(
			manners.TCPKeepAliveListener{tcpListener}, config)
	}

	s.srv = manners.NewWithOptions(
		manners.Options{
			Server:       s.newHTTPServer(),
			Listener:     listener,
			StateHandler: s.mux.connTracker.onStateChange,
		})
	s.state = srvStateHijacked
	return nil
}

func (s *srv) newHTTPServer() *http.Server {
	return &http.Server{
		Handler:        s.proxy,
		ReadTimeout:    s.options.ReadTimeout,
		WriteTimeout:   s.options.WriteTimeout,
		MaxHeaderBytes: s.options.MaxHeaderBytes,
	}
}

func (s *srv) reload() error {
	if !s.isServing() {
		return nil
	}

	// in case if the TLS in not served, we dont' need to do anything as it's all done by the proxy
	if !s.isTLS() {
		return nil
	}

	// Otherwise, we need to generate new TLS config and spin up the new server on the same socket
	config, err := newTLSConfig(s.keyPairs, s.defaultHost)
	if err != nil {
		return err
	}
	gracefulServer, err := s.srv.HijackListener(s.newHTTPServer(), config)
	if err != nil {
		return err
	}
	go s.serve(gracefulServer)

	s.srv.Close()
	s.srv = gracefulServer
	return nil
}

func (s *srv) shutdown() {
	if s.srv != nil {
		s.srv.Close()
	}
}

func newTLSConfig(keyPairs map[engine.HostKey]engine.KeyPair, defaultHost string) (*tls.Config, error) {
	config := &tls.Config{}

	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	pairs := make(map[string]tls.Certificate, len(keyPairs))
	for h, c := range keyPairs {
		keyPair, err := tls.X509KeyPair(c.Cert, c.Key)
		if err != nil {
			return nil, err
		}
		pairs[h.Name] = keyPair
	}

	config.Certificates = make([]tls.Certificate, 0, len(keyPairs))
	if defaultHost != "" {
		keyPair, exists := pairs[defaultHost]
		if !exists {
			return nil, fmt.Errorf("default host '%s' certificate is not passed", defaultHost)
		}
		config.Certificates = append(config.Certificates, keyPair)
	}

	for h, keyPair := range pairs {
		if h != defaultHost {
			config.Certificates = append(config.Certificates, keyPair)
		}
	}

	config.BuildNameToCertificate()
	return config, nil
}

func (s *srv) start() error {
	log.Infof("%s start", s)
	switch s.state {
	case srvStateInit:
		listener, err := net.Listen(s.listener.Address.Network, s.listener.Address.Address)
		if err != nil {
			return err
		}

		if s.isTLS() {
			config, err := newTLSConfig(s.keyPairs, s.defaultHost)
			if err != nil {
				return err
			}
			listener = manners.NewTLSListener(
				manners.TCPKeepAliveListener{listener.(*net.TCPListener)}, config)
		}
		s.srv = manners.NewWithOptions(
			manners.Options{
				Server:       s.newHTTPServer(),
				Listener:     listener,
				StateHandler: s.mux.connTracker.onStateChange,
			})
		s.state = srvStateActive
		go s.serve(s.srv)
		return nil
	case srvStateHijacked:
		s.state = srvStateActive
		go s.serve(s.srv)
		return nil
	}
	return fmt.Errorf("%v Calling start in unsupported state", s)
}

func (s *srv) serve(srv *manners.GracefulServer) {
	log.Infof("%s serve", s)

	s.mux.wg.Add(1)
	defer s.mux.wg.Done()

	srv.ListenAndServe()

	log.Infof("%v stop", s)
}

type srvState int

const (
	srvStateInit     = iota // server has been created
	srvStateActive   = iota // server is active and is serving requests
	srvStateHijacked = iota // server has hijacked listeners from other server
)

func (s srvState) String() string {
	switch s {
	case srvStateInit:
		return "init"
	case srvStateActive:
		return "active"
	case srvStateHijacked:
		return "hijacked"
	}
	return "undefined"
}
<commit_msg>Stronger crypto<commit_after>package proxy

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"

	"github.com/mailgun/vulcand/engine"

	"github.com/mailgun/vulcand/Godeps/_workspace/src/github.com/mailgun/log"
	"github.com/mailgun/vulcand/Godeps/_workspace/src/github.com/mailgun/manners"
)

// srv contains all that is necessary to run the HTTP(s) server. server does not work on its own,
// it heavily depends on MuxServer and acts as its internal data structure.
type srv struct {
	defaultHost string
	mux         *mux
	srv         *manners.GracefulServer
	proxy       http.Handler
	listener    engine.Listener
	keyPairs    map[engine.HostKey]engine.KeyPair
	options     Options
	state       int
}

func (s *srv) GetFile() (*FileDescriptor, error) {
	if !s.hasListeners() || s.srv == nil {
		return nil, nil
	}
	file, err := s.srv.GetFile()
	if err != nil {
		return nil, err
	}
	return &FileDescriptor{
		File:    file,
		Address: s.listener.Address,
	}, nil
}

func (s *srv) String() string {
	return fmt.Sprintf("%s->srv(%v, %v)", s.mux, s.state, &s.listener)
}

func newSrv(m *mux, l engine.Listener) (*srv, error) {
	defaultHost := ""
	keyPairs := make(map[engine.HostKey]engine.KeyPair)
	for hk, h := range m.hosts {
		if h.Settings.KeyPair != nil {
			keyPairs[hk] = *h.Settings.KeyPair
		}
		if h.Settings.Default {
			defaultHost = hk.Name
		}
	}

	return &srv{
		mux:         m,
		proxy:       m.router,
		listener:    l,
		defaultHost: defaultHost,
		keyPairs:    keyPairs,
		state:       srvStateInit,
	}, nil
}

func (s *srv) deleteKeyPair(hk engine.HostKey) error {
	delete(s.keyPairs, hk)
	return s.reload()
}

func (s *srv) isTLS() bool {
	return s.listener.Protocol == engine.HTTPS
}

func (s *srv) upsertKeyPair(hk engine.HostKey, keyPair *engine.KeyPair) error {
	old, exists := s.keyPairs[hk]
	if exists && old.Equals(keyPair) {
		return nil
	}
	s.keyPairs[hk] = *keyPair
	return s.reload()
}

func (s *srv) setDefaultHost(host engine.Host) error {
	oldDefault := s.defaultHost
	if host.Settings.Default {
		s.defaultHost = host.Name
	}
	if oldDefault != s.defaultHost && s.isTLS() {
		return s.reload()
	}
	return nil
}

func (s *srv) isServing() bool {
	return s.state == srvStateActive
}

func (s *srv) hasListeners() bool {
	return s.state == srvStateActive || s.state == srvStateHijacked
}

func (s *srv) takeFile(f *FileDescriptor) error {
	log.Infof("%s takeFile %v", s, f)

	listener, err := f.ToListener()
	if err != nil {
		return err
	}

	if s.isTLS() {
		tcpListener, ok := listener.(*net.TCPListener)
		if !ok {
			return fmt.Errorf(`%s failed to take file descriptor - it is running in TLS mode so I need a TCP listener,
but the file descriptor that was given corresponded to a listener of type %T. More about file descriptor: %s`, listener, s, f)
		}
		config, err := newTLSConfig(s.keyPairs, s.defaultHost)
		if err != nil {
			return err
		}
		listener = manners.NewTLSListener(
			manners.TCPKeepAliveListener{tcpListener}, config)
	}

	s.srv = manners.NewWithOptions(
		manners.Options{
			Server:       s.newHTTPServer(),
			Listener:     listener,
			StateHandler: s.mux.connTracker.onStateChange,
		})
	s.state = srvStateHijacked
	return nil
}

func (s *srv) newHTTPServer() *http.Server {
	return &http.Server{
		Handler:        s.proxy,
		ReadTimeout:    s.options.ReadTimeout,
		WriteTimeout:   s.options.WriteTimeout,
		MaxHeaderBytes: s.options.MaxHeaderBytes,
	}
}

func (s *srv) reload() error {
	if !s.isServing() {
		return nil
	}

	// in case if the TLS in not served, we dont' need to do anything as it's all done by the proxy
	if !s.isTLS() {
		return nil
	}

	// Otherwise, we need to generate new TLS config and spin up the new server on the same socket
	config, err := newTLSConfig(s.keyPairs, s.defaultHost)
	if err != nil {
		return err
	}
	gracefulServer, err := s.srv.HijackListener(s.newHTTPServer(), config)
	if err != nil {
		return err
	}
	go s.serve(gracefulServer)

	s.srv.Close()
	s.srv = gracefulServer
	return nil
}

func (s *srv) shutdown() {
	if s.srv != nil {
		s.srv.Close()
	}
}

func newTLSConfig(keyPairs map[engine.HostKey]engine.KeyPair, defaultHost string) (*tls.Config, error) {
	config := &tls.Config{}

	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	// only support TLS (mitigate against POODLE exploit)
	config.MinVersion = tls.VersionTLS10
	// use only modern ciphers
	config.CipherSuites = []uint16{
		tls.TLS_RSA_WITH_AES_128_CBC_SHA,
		tls.TLS_RSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
		tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
		tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
	}

	pairs := make(map[string]tls.Certificate, len(keyPairs))
	for h, c := range keyPairs {
		keyPair, err := tls.X509KeyPair(c.Cert, c.Key)
		if err != nil {
			return nil, err
		}
		pairs[h.Name] = keyPair
	}

	config.Certificates = make([]tls.Certificate, 0, len(keyPairs))
	if defaultHost != "" {
		keyPair, exists := pairs[defaultHost]
		if !exists {
			return nil, fmt.Errorf("default host '%s' certificate is not passed", defaultHost)
		}
		config.Certificates = append(config.Certificates, keyPair)
	}

	for h, keyPair := range pairs {
		if h != defaultHost {
			config.Certificates = append(config.Certificates, keyPair)
		}
	}

	config.BuildNameToCertificate()
	return config, nil
}

func (s *srv) start() error {
	log.Infof("%s start", s)
	switch s.state {
	case srvStateInit:
		listener, err := net.Listen(s.listener.Address.Network, s.listener.Address.Address)
		if err != nil {
			return err
		}

		if s.isTLS() {
			config, err := newTLSConfig(s.keyPairs, s.defaultHost)
			if err != nil {
				return err
			}
			listener = manners.NewTLSListener(
				manners.TCPKeepAliveListener{listener.(*net.TCPListener)}, config)
		}
		s.srv = manners.NewWithOptions(
			manners.Options{
				Server:       s.newHTTPServer(),
				Listener:     listener,
				StateHandler: s.mux.connTracker.onStateChange,
			})
		s.state = srvStateActive
		go s.serve(s.srv)
		return nil
	case srvStateHijacked:
		s.state = srvStateActive
		go s.serve(s.srv)
		return nil
	}
	return fmt.Errorf("%v Calling start in unsupported state", s)
}

func (s *srv) serve(srv *manners.GracefulServer) {
	log.Infof("%s serve", s)

	s.mux.wg.Add(1)
	defer s.mux.wg.Done()

	srv.ListenAndServe()

	log.Infof("%v stop", s)
}

type srvState int

const (
	srvStateInit     = iota // server has been created
	srvStateActive   = iota // server is active and is serving requests
	srvStateHijacked = iota // server has hijacked listeners from other server
)

func (s srvState) String() string {
	switch s {
	case srvStateInit:
		return "init"
	case srvStateActive:
		return "active"
	case srvStateHijacked:
		return "hijacked"
	}
	return "undefined"
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"
	"strings"
)

func main() {
	http.HandleFunc("/", handlerText)
	http.HandleFunc("/json", handlerJson)
	http.ListenAndServe(GetPort(), nil)
}

// Get the Port from the environment so we can run on Heroku
func GetPort() string {
	var port = os.Getenv("PORT")
	// Set a default port if there is nothing in the environment
	if port == "" {
		port = "8083"
		fmt.Println("INFO: No PORT environment variable detected, defaulting to " + port)
	}
	return ":" + port
}

func GetRemoteAddr(r *http.Request) string {
	ip := r.Header.Get("X-Forwarded-For")
	if len(ip) == 0 || strings.EqualFold("unknown", ip) {
		ip = strings.Split(r.RemoteAddr, ":")[0]
	}
	return ip
}

func handlerText(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, GetRemoteAddr(r))
}

func handlerJson(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "{ \"ip\": \"%s\" }", GetRemoteAddr(r))
}
<commit_msg>Set correct header when returning json.<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"
	"strings"
)

func main() {
	http.HandleFunc("/", handlerText)
	http.HandleFunc("/json", handlerJson)
	http.ListenAndServe(GetPort(), nil)
}

// Get the Port from the environment so we can run on Heroku
func GetPort() string {
	var port = os.Getenv("PORT")
	// Set a default port if there is nothing in the environment
	if port == "" {
		port = "8083"
		fmt.Println("INFO: No PORT environment variable detected, defaulting to " + port)
	}
	return ":" + port
}

func GetRemoteAddr(r *http.Request) string {
	ip := r.Header.Get("X-Forwarded-For")
	if len(ip) == 0 || strings.EqualFold("unknown", ip) {
		ip = strings.Split(r.RemoteAddr, ":")[0]
	}
	return ip
}

func handlerText(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, GetRemoteAddr(r))
}

func handlerJson(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-type", "application/json")
	fmt.Fprintf(w, "{ \"ip\": \"%s\" }", GetRemoteAddr(r))
}
<|endoftext|>


<commit_before>// Copyright (C) 2018 Google Inc.
//
// 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 validate

import (
	"context"
	"fmt"
	"math"

	"github.com/google/gapid/core/log"
	"github.com/google/gapid/gapis/perfetto"
	perfetto_service "github.com/google/gapid/gapis/perfetto/service"
)

const (
	counterIdQuery = "" +
		"select counter_id from counter_definitions " +
		"where name = '%v'"
	counterValuesQuery = "" +
		"select value from counter_values " +
		"where counter_id = %v order by ts " +
		"limit %v offset 10"
	renderStageTrackIdsQuery = "" +
		"select id from gpu_track " +
		"where scope = 'gpu_render_stage'"
	sampleCounter = 100
)

type Checker func(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool

type GpuCounter struct {
	Id    uint32
	Name  string
	Check Checker
}

type Validator interface {
	Validate(ctx context.Context, processor *perfetto.Processor) error
	GetCounters() []GpuCounter
}

func And(c1, c2 Checker) Checker {
	return func(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
		return c1(column, columnType) && c2(column, columnType)
	}
}

func IsNumber(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
	if columnType != perfetto_service.QueryResult_ColumnDesc_LONG && columnType != perfetto_service.QueryResult_ColumnDesc_DOUBLE {
		return false
	}
	return true
}

func CheckLargerThanZero(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
	longValues := column.GetLongValues()
	doubleValues := column.GetDoubleValues()
	for i := 0; i < sampleCounter; i++ {
		if columnType == perfetto_service.QueryResult_ColumnDesc_LONG {
			if longValues[i] <= 0 {
				return false
			}
		} else if columnType == perfetto_service.QueryResult_ColumnDesc_DOUBLE {
			if doubleValues[i] <= 0.0 {
				return false
			}
		}
	}
	return true
}

func CheckEqualTo(num float64) Checker {
	return func(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
		longValues := column.GetLongValues()
		doubleValues := column.GetDoubleValues()
		for i := 0; i < sampleCounter; i++ {
			if columnType == perfetto_service.QueryResult_ColumnDesc_LONG {
				if longValues[i] != int64(num) {
					return false
				}
			} else if columnType == perfetto_service.QueryResult_ColumnDesc_DOUBLE {
				if doubleValues[i] != num {
					return false
				}
			}
		}
		return true
	}
}

func CheckApproximateTo(num, err float64) Checker {
	return func(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
		longValues := column.GetLongValues()
		doubleValues := column.GetDoubleValues()
		for i := 0; i < sampleCounter; i++ {
			if columnType == perfetto_service.QueryResult_ColumnDesc_LONG {
				if math.Abs(num-float64(longValues[i])) > err {
					return false
				}
			} else if columnType == perfetto_service.QueryResult_ColumnDesc_DOUBLE {
				if math.Abs(num-doubleValues[i]) > err {
					return false
				}
			}
		}
		return true
	}
}

// GPU counters validation will fail in the below cases:
// 1. Fail to query
// 2. Missing GPU counter samples
// 3. Fail to check
func ValidateGpuCounters(ctx context.Context, processor *perfetto.Processor, counters []GpuCounter) error {
	for _, counter := range counters {
		queryResult, err := processor.Query(fmt.Sprintf(counterIdQuery, counter.Name))
		if err != nil {
			return log.Errf(ctx, err, "Failed to query with %v", fmt.Sprintf(counterIdQuery, counter.Name))
		}
		if len(queryResult.GetColumns()) != 1 {
			return log.Errf(ctx, err, "Expect one result with query: %v", fmt.Sprintf(counterIdQuery, counter.Name))
		}
		var counterId int64
		for _, column := range queryResult.GetColumns() {
			longValues := column.GetLongValues()
			if len(longValues) != 1 {
				// This should never happen, but sill have a check.
				return log.Err(ctx, nil, "Query result is not 1.")
			}
			counterId = longValues[0]
			break
		}
		queryResult, err = processor.Query(fmt.Sprintf(counterValuesQuery, counterId, sampleCounter))
		if err != nil {
			return log.Errf(ctx, err, "Failed to query with %v for counter %v", fmt.Sprintf(counterValuesQuery, counterId), counter)
		}

		// Query exactly #sampleCounter samples, fail if not enough samples
		if queryResult.GetNumRecords() != sampleCounter {
			return log.Errf(ctx, nil, "Number of samples is incorrect for counter: %v %v", counter, queryResult.GetNumRecords())
		}

		if !counter.Check(queryResult.GetColumns()[0], queryResult.GetColumnDescriptors()[0].GetType()) {
			return log.Errf(ctx, nil, "Check failed for counter: %v", counter)
		}
	}
	return nil
}

// GetRenderStageTrackIDs returns all track ids from gpu_track where the scope is gpu_render_stage
func GetRenderStageTrackIDs(ctx context.Context, processor *perfetto.Processor) ([]int64, error) {
	queryResult, err := processor.Query(renderStageTrackIdsQuery)
	if err != nil || queryResult.GetNumRecords() <= 0 {
		return []int64{}, log.Err(ctx, err, "Failed to query GPU render stage track ids")
	}
	result := make([]int64, queryResult.GetNumRecords())
	for i, v := range queryResult.GetColumns()[0].GetLongValues() {
		result[i] = v
	}
	return result, nil
}
<commit_msg>Update the queries for device validation.<commit_after>// Copyright (C) 2018 Google Inc.
//
// 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 validate

import (
	"context"
	"fmt"
	"math"

	"github.com/google/gapid/core/log"
	"github.com/google/gapid/gapis/perfetto"
	perfetto_service "github.com/google/gapid/gapis/perfetto/service"
)

const (
	counterIDQuery     = "select id from gpu_counter_track where name = '%v'"
	counterValuesQuery = "" +
		"select value from counter " +
		"where track_id = %v order by ts " +
		"limit %v offset 10"
	renderStageTrackIDQuery = "select id from gpu_track where scope = 'gpu_render_stage'"
	sampleCounter           = 100
)

// Checker is a function that checks the validity of the values of the given result set column.
type Checker func(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool

// GpuCounter represents a GPU counter for which the profiling data is validated.
type GpuCounter struct {
	Id    uint32
	Name  string
	Check Checker
}

// Validator is an interface implemented by the various hardware specific validators.
type Validator interface {
	Validate(ctx context.Context, processor *perfetto.Processor) error
	GetCounters() []GpuCounter
}

// And returns a checker that is only valid if both of its arguments are.
func And(c1, c2 Checker) Checker {
	return func(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
		return c1(column, columnType) && c2(column, columnType)
	}
}

// IsNumber is a checker that checks that the column is a number type.
func IsNumber(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
	if columnType != perfetto_service.QueryResult_ColumnDesc_LONG && columnType != perfetto_service.QueryResult_ColumnDesc_DOUBLE {
		return false
	}
	return true
}

// CheckLargerThanZero is a checker that checks that the values are all greater than zero.
func CheckLargerThanZero(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
	longValues := column.GetLongValues()
	doubleValues := column.GetDoubleValues()
	for i := 0; i < sampleCounter; i++ {
		if columnType == perfetto_service.QueryResult_ColumnDesc_LONG {
			if longValues[i] <= 0 {
				return false
			}
		} else if columnType == perfetto_service.QueryResult_ColumnDesc_DOUBLE {
			if doubleValues[i] <= 0.0 {
				return false
			}
		}
	}
	return true
}

// CheckEqualTo returns a checker that checks that all returned value equal the given value.
func CheckEqualTo(num float64) Checker {
	return func(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
		longValues := column.GetLongValues()
		doubleValues := column.GetDoubleValues()
		for i := 0; i < sampleCounter; i++ {
			if columnType == perfetto_service.QueryResult_ColumnDesc_LONG {
				if longValues[i] != int64(num) {
					return false
				}
			} else if columnType == perfetto_service.QueryResult_ColumnDesc_DOUBLE {
				if doubleValues[i] != num {
					return false
				}
			}
		}
		return true
	}
}

// CheckApproximateTo returns a checker that checks that values are within a margin of the given value.
func CheckApproximateTo(num, err float64) Checker {
	return func(column *perfetto_service.QueryResult_ColumnValues, columnType perfetto_service.QueryResult_ColumnDesc_Type) bool {
		longValues := column.GetLongValues()
		doubleValues := column.GetDoubleValues()
		for i := 0; i < sampleCounter; i++ {
			if columnType == perfetto_service.QueryResult_ColumnDesc_LONG {
				if math.Abs(num-float64(longValues[i])) > err {
					return false
				}
			} else if columnType == perfetto_service.QueryResult_ColumnDesc_DOUBLE {
				if math.Abs(num-doubleValues[i]) > err {
					return false
				}
			}
		}
		return true
	}
}

// ValidateGpuCounters validates the GPU counters.
// GPU counters validation will fail in the below cases:
// 1. Fail to query
// 2. Missing GPU counter samples
// 3. Fail to check
func ValidateGpuCounters(ctx context.Context, processor *perfetto.Processor, counters []GpuCounter) error {
	for _, counter := range counters {
		queryResult, err := processor.Query(fmt.Sprintf(counterIDQuery, counter.Name))
		if err != nil {
			return log.Errf(ctx, err, "Failed to query with %v", fmt.Sprintf(counterIDQuery, counter.Name))
		}
		if len(queryResult.GetColumns()) != 1 {
			return log.Errf(ctx, err, "Expect one result with query: %v", fmt.Sprintf(counterIDQuery, counter.Name))
		}
		var counterID int64
		for _, column := range queryResult.GetColumns() {
			longValues := column.GetLongValues()
			if len(longValues) != 1 {
				// This should never happen, but sill have a check.
				return log.Err(ctx, nil, "Query result is not 1.")
			}
			counterID = longValues[0]
			break
		}
		queryResult, err = processor.Query(fmt.Sprintf(counterValuesQuery, counterID, sampleCounter))
		if err != nil {
			return log.Errf(ctx, err, "Failed to query with %v for counter %v", fmt.Sprintf(counterValuesQuery, counterID), counter)
		}

		// Query exactly #sampleCounter samples, fail if not enough samples
		if queryResult.GetNumRecords() != sampleCounter {
			return log.Errf(ctx, nil, "Number of samples is incorrect for counter: %v %v", counter, queryResult.GetNumRecords())
		}

		if !counter.Check(queryResult.GetColumns()[0], queryResult.GetColumnDescriptors()[0].GetType()) {
			return log.Errf(ctx, nil, "Check failed for counter: %v", counter)
		}
	}
	return nil
}

// GetRenderStageTrackIDs returns all track ids from gpu_track where the scope is gpu_render_stage
func GetRenderStageTrackIDs(ctx context.Context, processor *perfetto.Processor) ([]int64, error) {
	queryResult, err := processor.Query(renderStageTrackIDQuery)
	if err != nil || queryResult.GetNumRecords() <= 0 {
		return []int64{}, log.Err(ctx, err, "Failed to query GPU render stage track ids")
	}
	result := make([]int64, queryResult.GetNumRecords())
	for i, v := range queryResult.GetColumns()[0].GetLongValues() {
		result[i] = v
	}
	return result, nil
}
<|endoftext|>


<commit_before>package rst

import (
	"net/http"
	"strconv"
	"strings"
	"time"
)

func normalizeHeaderArray(headers []string) []string {
	for i, name := range headers {
		headers[i] = http.CanonicalHeaderKey(name)
	}
	return headers
}

var defaultExposedHeaders = []string{http.CanonicalHeaderKey("etag")}

// DefaultAccessControl defines a limited CORS policy that only allows simple
// cross-origin requests.
var DefaultAccessControl = &AccessControlResponse{
	Origin:         "*",
	Credentials:    true,
	AllowedHeaders: nil,
	ExposedHeaders: defaultExposedHeaders,
	Methods:        nil,
	MaxAge:         24 * time.Hour,
}

// PermissiveAccessControl defines a permissive CORS policy in which all methods
// and all headers are allowed for all origins.
var PermissiveAccessControl = &AccessControlResponse{
	Origin:         "*",
	Credentials:    true,
	AllowedHeaders: []string{},
	ExposedHeaders: defaultExposedHeaders,
	Methods:        []string{},
	MaxAge:         24 * time.Hour,
}

/*
Preflighter is implemented by endpoints wishing to customize the response to
a CORS preflighted request.

	func (e *endpoint) Preflight(req *rst.AccessControlRequest, vars rst.RouteVars, r *http.Request) *rst.AccessControlResponse {
		if time.Now().Hour() < 12 {
			return &rst.AccessControlResponse{
				Origin: "morning.example.com",
				Methods: []string{"GET"},
			}
		}

		return &rst.AccessControlResponse{
			Origin: "afternoon.example.com",
			Methods: []string{"POST"},
		}
	}
*/
type Preflighter interface {
	Preflight(*AccessControlRequest, RouteVars, *http.Request) *AccessControlResponse
}

// AccessControlRequest represents the headers of a CORS access control request.
type AccessControlRequest struct {
	Origin  string
	Method  string
	Headers []string
}

func (ac *AccessControlRequest) isEmpty() bool {
	return ac.Origin == "" && ac.Method == "" && len(ac.Headers) == 0
}

// ParseAccessControlRequest returns a new instance of AccessControlRequest
// filled with CORS headers found in r.
func ParseAccessControlRequest(r *http.Request) *AccessControlRequest {
	var headers []string
	if h := r.Header.Get("Access-Control-Request-Headers"); h != "" {
		headers = strings.Split(strings.Replace(r.Header.Get("Access-Control-Request-Headers"), " ", "", -1), ",")
	}
	return &AccessControlRequest{
		Origin:  r.Header.Get("Origin"),
		Method:  r.Header.Get("Access-Control-Request-Method"),
		Headers: headers,
	}

	// TODO: remove duplicated headers before serving them back.
}

// AccessControlResponse defines the response headers to a CORS access control
// request.
type AccessControlResponse struct {
	Origin         string
	ExposedHeaders []string
	Methods        []string // Empty array means any, nil means none.
	AllowedHeaders []string // Empty array means any, nil means none.
	Credentials    bool
	MaxAge         time.Duration
}

type accessControlHandler struct {
	endpoint Endpoint
	*AccessControlResponse
}

func newAccessControlHandler(endpoint Endpoint, ac *AccessControlResponse) *accessControlHandler {
	return &accessControlHandler{
		endpoint:              endpoint,
		AccessControlResponse: ac,
	}
}

func (h *accessControlHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if _, exists := r.Header["Origin"]; !exists {
		return
	}

	req := ParseAccessControlRequest(r)

	var resp *AccessControlResponse
	if h.endpoint == nil {
		resp = h.AccessControlResponse
	} else {
		if preflighter, implemented := h.endpoint.(Preflighter); implemented && strings.ToUpper(r.Method) == Options {
			// If Options and endpoint implements Preflighter, call Preflight.
			resp = preflighter.Preflight(req, getVars(r), r)
		} else {
			resp = h.AccessControlResponse
		}
	}

	// Adding a vary if an origin is specified in the response.
	defer func() {
		if allowed := w.Header().Get("Access-Control-Allow-Origin"); allowed != "" && allowed != "*" {
			w.Header().Add("Vary", "Origin")
		}
	}()

	// Writing response headers
	if resp.Origin != "" {
		w.Header().Set("Access-Control-Allow-Origin", resp.Origin)
	}
	w.Header().Set("Access-Control-Allow-Credentials", strconv.FormatBool(resp.Credentials))

	// Exposed headers
	if len(resp.ExposedHeaders) > 0 {
		w.Header().Set("Access-Control-Expose-Headers", strings.Join(normalizeHeaderArray(resp.ExposedHeaders), ", "))
	}

	// OPTIONS only
	if strings.ToUpper(r.Method) != Options {
		return
	}

	w.Header().Set("Access-Control-Max-Age", strconv.Itoa(int(resp.MaxAge.Seconds())))

	if req.Method != "" && resp.Methods != nil {
		var methods []string
		if len(resp.Methods) == 0 {
			methods = AllowedMethods(h.endpoint)
		} else {
			methods = resp.Methods
		}
		w.Header().Set("Access-Control-Allow-Methods", strings.Join(methods, ", "))
	}

	if len(req.Headers) > 0 && resp.AllowedHeaders != nil {
		var headers []string
		if len(resp.AllowedHeaders) == 0 {
			headers = req.Headers
		} else {
			headers = resp.AllowedHeaders
		}
		w.Header().Set("Access-Control-Allow-Headers", strings.Join(normalizeHeaderArray(headers), ", "))
	}
}
<commit_msg>Changed visibility of the default exposed headers<commit_after>package rst

import (
	"net/http"
	"strconv"
	"strings"
	"time"
)

func normalizeHeaderArray(headers []string) []string {
	for i, name := range headers {
		headers[i] = http.CanonicalHeaderKey(name)
	}
	return headers
}

// DefaultAccessControl defines a limited CORS policy that only allows simple
// cross-origin requests.
var DefaultAccessControl = &AccessControlResponse{
	Origin:         "*",
	Credentials:    true,
	AllowedHeaders: nil,
	ExposedHeaders: []string{"Etag"),
	Methods:        nil,
	MaxAge:         24 * time.Hour,
}

// PermissiveAccessControl defines a permissive CORS policy in which all methods
// and all headers are allowed for all origins.
var PermissiveAccessControl = &AccessControlResponse{
	Origin:         "*",
	Credentials:    true,
	AllowedHeaders: []string{},
	ExposedHeaders: []string{"Etag"),
	Methods:        []string{},
	MaxAge:         24 * time.Hour,
}

/*
Preflighter is implemented by endpoints wishing to customize the response to
a CORS preflighted request.

	func (e *endpoint) Preflight(req *rst.AccessControlRequest, vars rst.RouteVars, r *http.Request) *rst.AccessControlResponse {
		if time.Now().Hour() < 12 {
			return &rst.AccessControlResponse{
				Origin: "morning.example.com",
				Methods: []string{"GET"},
			}
		}

		return &rst.AccessControlResponse{
			Origin: "afternoon.example.com",
			Methods: []string{"POST"},
		}
	}
*/
type Preflighter interface {
	Preflight(*AccessControlRequest, RouteVars, *http.Request) *AccessControlResponse
}

// AccessControlRequest represents the headers of a CORS access control request.
type AccessControlRequest struct {
	Origin  string
	Method  string
	Headers []string
}

func (ac *AccessControlRequest) isEmpty() bool {
	return ac.Origin == "" && ac.Method == "" && len(ac.Headers) == 0
}

// ParseAccessControlRequest returns a new instance of AccessControlRequest
// filled with CORS headers found in r.
func ParseAccessControlRequest(r *http.Request) *AccessControlRequest {
	var headers []string
	if h := r.Header.Get("Access-Control-Request-Headers"); h != "" {
		headers = strings.Split(strings.Replace(r.Header.Get("Access-Control-Request-Headers"), " ", "", -1), ",")
	}
	return &AccessControlRequest{
		Origin:  r.Header.Get("Origin"),
		Method:  r.Header.Get("Access-Control-Request-Method"),
		Headers: headers,
	}

	// TODO: remove duplicated headers before serving them back.
}

// AccessControlResponse defines the response headers to a CORS access control
// request.
type AccessControlResponse struct {
	Origin         string
	ExposedHeaders []string
	Methods        []string // Empty array means any, nil means none.
	AllowedHeaders []string // Empty array means any, nil means none.
	Credentials    bool
	MaxAge         time.Duration
}

type accessControlHandler struct {
	endpoint Endpoint
	*AccessControlResponse
}

func newAccessControlHandler(endpoint Endpoint, ac *AccessControlResponse) *accessControlHandler {
	return &accessControlHandler{
		endpoint:              endpoint,
		AccessControlResponse: ac,
	}
}

func (h *accessControlHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if _, exists := r.Header["Origin"]; !exists {
		return
	}

	req := ParseAccessControlRequest(r)

	var resp *AccessControlResponse
	if h.endpoint == nil {
		resp = h.AccessControlResponse
	} else {
		if preflighter, implemented := h.endpoint.(Preflighter); implemented && strings.ToUpper(r.Method) == Options {
			// If Options and endpoint implements Preflighter, call Preflight.
			resp = preflighter.Preflight(req, getVars(r), r)
		} else {
			resp = h.AccessControlResponse
		}
	}

	// Adding a vary if an origin is specified in the response.
	defer func() {
		if allowed := w.Header().Get("Access-Control-Allow-Origin"); allowed != "" && allowed != "*" {
			w.Header().Add("Vary", "Origin")
		}
	}()

	// Writing response headers
	if resp.Origin != "" {
		w.Header().Set("Access-Control-Allow-Origin", resp.Origin)
	}
	w.Header().Set("Access-Control-Allow-Credentials", strconv.FormatBool(resp.Credentials))

	// Exposed headers
	if len(resp.ExposedHeaders) > 0 {
		w.Header().Set("Access-Control-Expose-Headers", strings.Join(normalizeHeaderArray(resp.ExposedHeaders), ", "))
	}

	// OPTIONS only
	if strings.ToUpper(r.Method) != Options {
		return
	}

	w.Header().Set("Access-Control-Max-Age", strconv.Itoa(int(resp.MaxAge.Seconds())))

	if req.Method != "" && resp.Methods != nil {
		var methods []string
		if len(resp.Methods) == 0 {
			methods = AllowedMethods(h.endpoint)
		} else {
			methods = resp.Methods
		}
		w.Header().Set("Access-Control-Allow-Methods", strings.Join(methods, ", "))
	}

	if len(req.Headers) > 0 && resp.AllowedHeaders != nil {
		var headers []string
		if len(resp.AllowedHeaders) == 0 {
			headers = req.Headers
		} else {
			headers = resp.AllowedHeaders
		}
		w.Header().Set("Access-Control-Allow-Headers", strings.Join(normalizeHeaderArray(headers), ", "))
	}
}
<|endoftext|>