<commit_before>package version

import (
	"os"
	"strconv"
)

// APIVersion contains the API base version. Only bumped for backward incompatible changes.
var APIVersion = "1.0"

// APIExtensions is the list of all API extensions in the order they were added.
//
// The following kind of changes come with a new extensions:
//
// - New configuration key
// - New valid values for a configuration key
// - New REST API endpoint
// - New argument inside an existing REST API call
// - New HTTPs authentication mechanisms or protocols
//
// This list is used mainly by the LXD server code, but it's in the shared
// package as well for reference.
var APIExtensions = []string{
	"storage_zfs_remove_snapshots",
	"container_host_shutdown_timeout",
	"container_stop_priority",
	"container_syscall_filtering",
	"auth_pki",
	"container_last_used_at",
	"etag",
	"patch",
	"usb_devices",
	"https_allowed_credentials",
	"image_compression_algorithm",
	"directory_manipulation",
	"container_cpu_time",
	"storage_zfs_use_refquota",
	"storage_lvm_mount_options",
	"network",
	"profile_usedby",
	"container_push",
	"container_exec_recording",
	"certificate_update",
	"container_exec_signal_handling",
	"gpu_devices",
	"container_image_properties",
	"migration_progress",
	"id_map",
	"network_firewall_filtering",
	"network_routes",
	"storage",
	"file_delete",
	"file_append",
	"network_dhcp_expiry",
	"storage_lvm_vg_rename",
	"storage_lvm_thinpool_rename",
	"network_vlan",
	"image_create_aliases",
	"container_stateless_copy",
	"container_only_migration",
	"storage_zfs_clone_copy",
	"unix_device_rename",
	"storage_lvm_use_thinpool",
	"storage_rsync_bwlimit",
	"network_vxlan_interface",
	"storage_btrfs_mount_options",
	"entity_description",
	"image_force_refresh",
	"storage_lvm_lv_resizing",
	"id_map_base",
	"file_symlinks",
	"container_push_target",
	"network_vlan_physical",
	"storage_images_delete",
	"container_edit_metadata",
	"container_snapshot_stateful_migration",
	"storage_driver_ceph",
	"storage_ceph_user_name",
	"resource_limits",
	"storage_volatile_initial_source",
	"storage_ceph_force_osd_reuse",
	"storage_block_filesystem_btrfs",
	"resources",
	"kernel_limits",
	"storage_api_volume_rename",
	"macaroon_authentication",
	"network_sriov",
	"console",
	"restrict_devlxd",
	"migration_pre_copy",
	"infiniband",
	"maas_network",
	"devlxd_events",
	"proxy",
	"network_dhcp_gateway",
	"file_get_symlink",
	"network_leases",
	"unix_device_hotplug",
	"storage_api_local_volume_handling",
	"operation_description",
	"clustering",
	"event_lifecycle",
	"storage_api_remote_volume_handling",
	"nvidia_runtime",
	"container_mount_propagation",
	"container_backup",
	"devlxd_images",
	"container_local_cross_pool_handling",
	"proxy_unix",
	"proxy_udp",
	"clustering_join",
	"proxy_tcp_udp_multi_port_handling",
	"network_state",
	"proxy_unix_dac_properties",
	"container_protection_delete",
	"unix_priv_drop",
	"pprof_http",
	"proxy_haproxy_protocol",
	"network_hwaddr",
	"proxy_nat",
	"network_nat_order",
	"container_full",
	"candid_authentication",
	"backup_compression",
	"candid_config",
	"nvidia_runtime_config",
	"storage_api_volume_snapshots",
	"storage_unmapped",
	"projects",
	"candid_config_key",
	"network_vxlan_ttl",
	"container_incremental_copy",
	"usb_optional_vendorid",
	"snapshot_scheduling",
	"container_copy_project",
	"clustering_server_address",
	"clustering_image_replication",
	"container_protection_shift",
	"snapshot_expiry",
	"container_backup_override_pool",
	"snapshot_expiry_creation",
	"network_leases_location",
	"resources_cpu_socket",
	"resources_gpu",
	"resources_numa",
	"kernel_features",
	"id_map_current",
	"event_location",
	"storage_api_remote_volume_snapshots",
	"network_nat_address",
	"container_nic_routes",
	"rbac",
	"cluster_internal_copy",
	"seccomp_notify",
	"lxc_features",
	"container_nic_ipvlan",
	"network_vlan_sriov",
	"storage_cephfs",
	"container_nic_ipfilter",
	"resources_v2",
	"container_exec_user_group_cwd",
	"container_syscall_intercept",
	"container_disk_shift",
	"storage_shifted",
	"resources_infiniband",
	"daemon_storage",
	"instances",
	"image_types",
	"resources_disk_sata",
	"clustering_roles",
	"images_expiry",
	"resources_network_firmware",
	"backup_compression_algorithm",
	"ceph_data_pool_name",
	"container_syscall_intercept_mount",
	"compression_squashfs",
	"container_raw_mount",
	"container_nic_routed",
	"container_syscall_intercept_mount_fuse",
	"container_disk_ceph",
	"virtual-machines",
	"image_profiles",
	"clustering_architecture",
	"resources_disk_id",
	"storage_lvm_stripes",
	"vm_boot_priority",
	"unix_hotplug_devices",
	"api_filtering",
	"instance_nic_network",
	"clustering_sizing",
	"firewall_driver",
	"projects_limits",
	"container_syscall_intercept_hugetlbfs",
	"limits_hugepages",
	"container_nic_routed_gateway",
	"projects_restrictions",
}

// APIExtensionsCount returns the number of available API extensions.
func APIExtensionsCount() int {
	count := len(APIExtensions)

	// This environment variable is an internal one to force the code
	// to believe that we have an API extensions count greater than we
	// actually have. It's used by integration tests to exercise the
	// cluster upgrade process.
	artificialBump := os.Getenv("LXD_ARTIFICIALLY_BUMP_API_EXTENSIONS")
	if artificialBump != "" {
		n, err := strconv.Atoi(artificialBump)
		if err == nil {
			count += n
		}
	}

	return count
}
<commit_msg>shared/version/api: Add custom_volume_snapshot_expiry extension<commit_after>package version

import (
	"os"
	"strconv"
)

// APIVersion contains the API base version. Only bumped for backward incompatible changes.
var APIVersion = "1.0"

// APIExtensions is the list of all API extensions in the order they were added.
//
// The following kind of changes come with a new extensions:
//
// - New configuration key
// - New valid values for a configuration key
// - New REST API endpoint
// - New argument inside an existing REST API call
// - New HTTPs authentication mechanisms or protocols
//
// This list is used mainly by the LXD server code, but it's in the shared
// package as well for reference.
var APIExtensions = []string{
	"storage_zfs_remove_snapshots",
	"container_host_shutdown_timeout",
	"container_stop_priority",
	"container_syscall_filtering",
	"auth_pki",
	"container_last_used_at",
	"etag",
	"patch",
	"usb_devices",
	"https_allowed_credentials",
	"image_compression_algorithm",
	"directory_manipulation",
	"container_cpu_time",
	"storage_zfs_use_refquota",
	"storage_lvm_mount_options",
	"network",
	"profile_usedby",
	"container_push",
	"container_exec_recording",
	"certificate_update",
	"container_exec_signal_handling",
	"gpu_devices",
	"container_image_properties",
	"migration_progress",
	"id_map",
	"network_firewall_filtering",
	"network_routes",
	"storage",
	"file_delete",
	"file_append",
	"network_dhcp_expiry",
	"storage_lvm_vg_rename",
	"storage_lvm_thinpool_rename",
	"network_vlan",
	"image_create_aliases",
	"container_stateless_copy",
	"container_only_migration",
	"storage_zfs_clone_copy",
	"unix_device_rename",
	"storage_lvm_use_thinpool",
	"storage_rsync_bwlimit",
	"network_vxlan_interface",
	"storage_btrfs_mount_options",
	"entity_description",
	"image_force_refresh",
	"storage_lvm_lv_resizing",
	"id_map_base",
	"file_symlinks",
	"container_push_target",
	"network_vlan_physical",
	"storage_images_delete",
	"container_edit_metadata",
	"container_snapshot_stateful_migration",
	"storage_driver_ceph",
	"storage_ceph_user_name",
	"resource_limits",
	"storage_volatile_initial_source",
	"storage_ceph_force_osd_reuse",
	"storage_block_filesystem_btrfs",
	"resources",
	"kernel_limits",
	"storage_api_volume_rename",
	"macaroon_authentication",
	"network_sriov",
	"console",
	"restrict_devlxd",
	"migration_pre_copy",
	"infiniband",
	"maas_network",
	"devlxd_events",
	"proxy",
	"network_dhcp_gateway",
	"file_get_symlink",
	"network_leases",
	"unix_device_hotplug",
	"storage_api_local_volume_handling",
	"operation_description",
	"clustering",
	"event_lifecycle",
	"storage_api_remote_volume_handling",
	"nvidia_runtime",
	"container_mount_propagation",
	"container_backup",
	"devlxd_images",
	"container_local_cross_pool_handling",
	"proxy_unix",
	"proxy_udp",
	"clustering_join",
	"proxy_tcp_udp_multi_port_handling",
	"network_state",
	"proxy_unix_dac_properties",
	"container_protection_delete",
	"unix_priv_drop",
	"pprof_http",
	"proxy_haproxy_protocol",
	"network_hwaddr",
	"proxy_nat",
	"network_nat_order",
	"container_full",
	"candid_authentication",
	"backup_compression",
	"candid_config",
	"nvidia_runtime_config",
	"storage_api_volume_snapshots",
	"storage_unmapped",
	"projects",
	"candid_config_key",
	"network_vxlan_ttl",
	"container_incremental_copy",
	"usb_optional_vendorid",
	"snapshot_scheduling",
	"container_copy_project",
	"clustering_server_address",
	"clustering_image_replication",
	"container_protection_shift",
	"snapshot_expiry",
	"container_backup_override_pool",
	"snapshot_expiry_creation",
	"network_leases_location",
	"resources_cpu_socket",
	"resources_gpu",
	"resources_numa",
	"kernel_features",
	"id_map_current",
	"event_location",
	"storage_api_remote_volume_snapshots",
	"network_nat_address",
	"container_nic_routes",
	"rbac",
	"cluster_internal_copy",
	"seccomp_notify",
	"lxc_features",
	"container_nic_ipvlan",
	"network_vlan_sriov",
	"storage_cephfs",
	"container_nic_ipfilter",
	"resources_v2",
	"container_exec_user_group_cwd",
	"container_syscall_intercept",
	"container_disk_shift",
	"storage_shifted",
	"resources_infiniband",
	"daemon_storage",
	"instances",
	"image_types",
	"resources_disk_sata",
	"clustering_roles",
	"images_expiry",
	"resources_network_firmware",
	"backup_compression_algorithm",
	"ceph_data_pool_name",
	"container_syscall_intercept_mount",
	"compression_squashfs",
	"container_raw_mount",
	"container_nic_routed",
	"container_syscall_intercept_mount_fuse",
	"container_disk_ceph",
	"virtual-machines",
	"image_profiles",
	"clustering_architecture",
	"resources_disk_id",
	"storage_lvm_stripes",
	"vm_boot_priority",
	"unix_hotplug_devices",
	"api_filtering",
	"instance_nic_network",
	"clustering_sizing",
	"firewall_driver",
	"projects_limits",
	"container_syscall_intercept_hugetlbfs",
	"limits_hugepages",
	"container_nic_routed_gateway",
	"projects_restrictions",
	"custom_volume_snapshot_expiry",
}

// APIExtensionsCount returns the number of available API extensions.
func APIExtensionsCount() int {
	count := len(APIExtensions)

	// This environment variable is an internal one to force the code
	// to believe that we have an API extensions count greater than we
	// actually have. It's used by integration tests to exercise the
	// cluster upgrade process.
	artificialBump := os.Getenv("LXD_ARTIFICIALLY_BUMP_API_EXTENSIONS")
	if artificialBump != "" {
		n, err := strconv.Atoi(artificialBump)
		if err == nil {
			count += n
		}
	}

	return count
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// 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 sources

import (
	"fmt"
	"net"
	"net/http"
	"net/url"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	kube_api "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kube_client "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	kube_labels "github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/golang/glog"
	cadvisor "github.com/google/cadvisor/info"
)

// Kubernetes released supported and tested against.
var kubeVersions = []string{"v0.3"}

const (
	// Cadvisor port in kubernetes.
	cadvisorPort = 4194

	kubeClientVersion = "v1beta1"
)

type PodInstance struct {
	Pod    string
	PodId  string
	HostIp string
}

type KubeSource struct {
	client      *kube_client.Client
	lastQuery   time.Time
	kubeletPort string
	stateLock   sync.RWMutex
	goodNodes   []string            // guarded by stateLock
	nodeErrors  map[string]int      // guarded by stateLock
	podErrors   map[PodInstance]int // guarded by stateLock
}

type nodeList CadvisorHosts

func (self *KubeSource) recordNodeError(name string) {
	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	self.nodeErrors[name]++
}

func (self *KubeSource) recordPodError(pod Pod) {
	// Heapster knows about pods before they are up and running on a node.
	// Ignore errors for Pods that are not Running.
	if pod.Status != "Running" {
		return
	}

	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	podInstance := PodInstance{Pod: pod.Name, PodId: pod.ID, HostIp: pod.HostIP}
	self.podErrors[podInstance]++
}

func (self *KubeSource) recordGoodNodes(nodes []string) {
	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	self.goodNodes = nodes
}

func (self *KubeSource) getState() string {
	self.stateLock.RLock()
	defer self.stateLock.RUnlock()

	state := "\tHealthy Nodes:\n"
	for _, node := range self.goodNodes {
		state += fmt.Sprintf("\t\t%s\n", node)
	}
	if len(self.nodeErrors) != 0 {
		state += fmt.Sprintf("\tNode Errors: %+v\n", self.nodeErrors)
	} else {
		state += "\tNo node errors\n"
	}
	if len(self.podErrors) != 0 {
		state += fmt.Sprintf("\tPod Errors: %+v\n", self.podErrors)
	} else {
		state += "\tNo pod errors\n"
	}
	return state
}

// Returns a map of minion hostnames to their corresponding IPs.
func (self *KubeSource) listMinions() (*nodeList, error) {
	nodeList := &nodeList{
		Port:  cadvisorPort,
		Hosts: make(map[string]string, 0),
	}
	minions, err := self.client.Nodes().List()
	if err != nil {
		return nil, err
	}
	goodNodes := []string{}
	for _, minion := range minions.Items {
		addrs, err := net.LookupIP(minion.Name)
		if err == nil {
			nodeList.Hosts[minion.Name] = addrs[0].String()
			goodNodes = append(goodNodes, minion.Name)
		} else {
			glog.Errorf("Skipping host %s since looking up its IP failed - %s", minion.Name, err)
			self.recordNodeError(minion.Name)
		}
	}
	self.recordGoodNodes(goodNodes)

	return nodeList, nil
}

func (self *KubeSource) parsePod(pod *kube_api.Pod) *Pod {
	localPod := Pod{
		Name:       pod.Name,
		Namespace:  pod.Namespace,
		ID:         string(pod.UID),
		Hostname:   pod.Status.Host,
		Status:     string(pod.Status.Phase),
		PodIP:      pod.Status.PodIP,
		Labels:     make(map[string]string, 0),
		Containers: make([]*Container, 0),
	}
	for key, value := range pod.Labels {
		localPod.Labels[key] = value
	}
	for _, container := range pod.Spec.Containers {
		localContainer := newContainer()
		localContainer.Name = container.Name
		localPod.Containers = append(localPod.Containers, localContainer)
	}
	glog.V(2).Infof("found pod: %+v", localPod)

	return &localPod
}

// Returns a map of minion hostnames to the Pods running in them.
func (self *KubeSource) getPods() ([]Pod, error) {
	pods, err := self.client.Pods(kube_api.NamespaceAll).List(kube_labels.Everything())
	if err != nil {
		return nil, err
	}
	glog.V(1).Infof("got pods from api server %+v", pods)
	// TODO(vishh): Add API Version check. Fail if Kubernetes returns an invalid API Version.
	out := make([]Pod, 0)
	for _, pod := range pods.Items {
		glog.V(2).Infof("Got Kube Pod: %+v", pod)
		pod := self.parsePod(&pod)
		addrs, err := net.LookupIP(pod.Hostname)
		if err != nil {
			glog.Errorf("Skipping host %s since looking up its IP failed - %s", pod.Hostname, err)
			self.recordNodeError(pod.Hostname)
			continue
		}
		pod.HostIP = addrs[0].String()
		out = append(out, *pod)
	}

	return out, nil
}

func (self *KubeSource) getStatsFromKubelet(pod Pod, containerName string) (cadvisor.ContainerSpec, []*cadvisor.ContainerStats, error) {
	var containerInfo cadvisor.ContainerInfo
	values := url.Values{}
	values.Add("num_stats", strconv.Itoa(int(time.Since(self.lastQuery)/time.Second)))
	url := "http://" + pod.HostIP + ":" + self.kubeletPort + filepath.Join("/stats", pod.Namespace, pod.Name, pod.ID, containerName) + "?" + values.Encode()
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return cadvisor.ContainerSpec{}, []*cadvisor.ContainerStats{}, err
	}
	err = PostRequestAndGetValue(&http.Client{}, req, &containerInfo)
	if err != nil {
		glog.Errorf("failed to get stats from kubelet url: %s - %s\n", url, err)
		self.recordPodError(pod)
		return cadvisor.ContainerSpec{}, []*cadvisor.ContainerStats{}, nil
	}

	return containerInfo.Spec, containerInfo.Stats, nil
}

func (self *KubeSource) getNodesInfo() ([]RawContainer, error) {
	kubeNodes, err := self.listMinions()
	if err != nil {
		return []RawContainer{}, err
	}
	nodesInfo := []RawContainer{}
	for node, ip := range kubeNodes.Hosts {
		spec, stats, err := self.getStatsFromKubelet(Pod{HostIP: ip}, "/")
		if err != nil {
			glog.V(1).Infof("Failed to get machine stats from kubelet for node %s", node)
			return []RawContainer{}, err
		}
		if len(stats) > 0 {
			container := RawContainer{node, Container{"/", spec, stats}}
			nodesInfo = append(nodesInfo, container)
		}
	}

	return nodesInfo, nil
}

func (self *KubeSource) GetInfo() (ContainerData, error) {
	pods, err := self.getPods()
	if err != nil {
		return ContainerData{}, err
	}
	for _, pod := range pods {
		for _, container := range pod.Containers {
			spec, stats, err := self.getStatsFromKubelet(pod, container.Name)
			if err != nil {
				return ContainerData{}, err
			}
			glog.V(2).Infof("Fetched stats from kubelet for container %s in pod %s", container.Name, pod.Name)
			container.Stats = stats
			container.Spec = spec
		}
	}
	nodesInfo, err := self.getNodesInfo()
	if err != nil {
		return ContainerData{}, err
	}
	glog.V(2).Info("Fetched list of nodes from the master")
	self.lastQuery = time.Now()

	return ContainerData{Pods: pods, Machine: nodesInfo}, nil
}

func newKubeSource() (*KubeSource, error) {
	if len(*argMaster) == 0 {
		return nil, fmt.Errorf("kubernetes_master flag not specified")
	}
	kubeClient := kube_client.NewOrDie(&kube_client.Config{
		Host:     "http://" + *argMaster,
		Version:  kubeClientVersion,
		Insecure: true,
	})

	glog.Infof("Using Kubernetes client with master %q and version %s\n", *argMaster, kubeClientVersion)
	glog.Infof("Using kubelet port %q", *argKubeletPort)
	glog.Infof("Support kubelet versions %v", kubeVersions)

	return &KubeSource{
		client:      kubeClient,
		lastQuery:   time.Now(),
		kubeletPort: *argKubeletPort,
		nodeErrors:  make(map[string]int),
		podErrors:   make(map[PodInstance]int),
	}, nil
}

func (self *KubeSource) GetConfig() string {
	desc := "Source type: Kube\n"
	desc += fmt.Sprintf("\tClient config: master ip %q, version %s\n", *argMaster, kubeClientVersion)
	desc += fmt.Sprintf("\tUsing kubelet port %q\n", self.kubeletPort)
	desc += fmt.Sprintf("\tSupported kubelet versions %v\n", kubeVersions)
	desc += self.getState()
	desc += "\n"
	return desc
}
<commit_msg>Remove kube versions supported log message as it is outdated and not necessary since heapster versions are tagged to specific kubernetes versions.<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// 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 sources

import (
	"fmt"
	"net"
	"net/http"
	"net/url"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	kube_api "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kube_client "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	kube_labels "github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/golang/glog"
	cadvisor "github.com/google/cadvisor/info"
)

const (
	// Cadvisor port in kubernetes.
	cadvisorPort = 4194

	kubeClientVersion = "v1beta1"
)

type PodInstance struct {
	Pod    string
	PodId  string
	HostIp string
}

type KubeSource struct {
	client      *kube_client.Client
	lastQuery   time.Time
	kubeletPort string
	stateLock   sync.RWMutex
	goodNodes   []string            // guarded by stateLock
	nodeErrors  map[string]int      // guarded by stateLock
	podErrors   map[PodInstance]int // guarded by stateLock
}

type nodeList CadvisorHosts

func (self *KubeSource) recordNodeError(name string) {
	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	self.nodeErrors[name]++
}

func (self *KubeSource) recordPodError(pod Pod) {
	// Heapster knows about pods before they are up and running on a node.
	// Ignore errors for Pods that are not Running.
	if pod.Status != "Running" {
		return
	}

	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	podInstance := PodInstance{Pod: pod.Name, PodId: pod.ID, HostIp: pod.HostIP}
	self.podErrors[podInstance]++
}

func (self *KubeSource) recordGoodNodes(nodes []string) {
	self.stateLock.Lock()
	defer self.stateLock.Unlock()

	self.goodNodes = nodes
}

func (self *KubeSource) getState() string {
	self.stateLock.RLock()
	defer self.stateLock.RUnlock()

	state := "\tHealthy Nodes:\n"
	for _, node := range self.goodNodes {
		state += fmt.Sprintf("\t\t%s\n", node)
	}
	if len(self.nodeErrors) != 0 {
		state += fmt.Sprintf("\tNode Errors: %+v\n", self.nodeErrors)
	} else {
		state += "\tNo node errors\n"
	}
	if len(self.podErrors) != 0 {
		state += fmt.Sprintf("\tPod Errors: %+v\n", self.podErrors)
	} else {
		state += "\tNo pod errors\n"
	}
	return state
}

// Returns a map of minion hostnames to their corresponding IPs.
func (self *KubeSource) listMinions() (*nodeList, error) {
	nodeList := &nodeList{
		Port:  cadvisorPort,
		Hosts: make(map[string]string, 0),
	}
	minions, err := self.client.Nodes().List()
	if err != nil {
		return nil, err
	}
	goodNodes := []string{}
	for _, minion := range minions.Items {
		addrs, err := net.LookupIP(minion.Name)
		if err == nil {
			nodeList.Hosts[minion.Name] = addrs[0].String()
			goodNodes = append(goodNodes, minion.Name)
		} else {
			glog.Errorf("Skipping host %s since looking up its IP failed - %s", minion.Name, err)
			self.recordNodeError(minion.Name)
		}
	}
	self.recordGoodNodes(goodNodes)

	return nodeList, nil
}

func (self *KubeSource) parsePod(pod *kube_api.Pod) *Pod {
	localPod := Pod{
		Name:       pod.Name,
		Namespace:  pod.Namespace,
		ID:         string(pod.UID),
		Hostname:   pod.Status.Host,
		Status:     string(pod.Status.Phase),
		PodIP:      pod.Status.PodIP,
		Labels:     make(map[string]string, 0),
		Containers: make([]*Container, 0),
	}
	for key, value := range pod.Labels {
		localPod.Labels[key] = value
	}
	for _, container := range pod.Spec.Containers {
		localContainer := newContainer()
		localContainer.Name = container.Name
		localPod.Containers = append(localPod.Containers, localContainer)
	}
	glog.V(2).Infof("found pod: %+v", localPod)

	return &localPod
}

// Returns a map of minion hostnames to the Pods running in them.
func (self *KubeSource) getPods() ([]Pod, error) {
	pods, err := self.client.Pods(kube_api.NamespaceAll).List(kube_labels.Everything())
	if err != nil {
		return nil, err
	}
	glog.V(1).Infof("got pods from api server %+v", pods)
	// TODO(vishh): Add API Version check. Fail if Kubernetes returns an invalid API Version.
	out := make([]Pod, 0)
	for _, pod := range pods.Items {
		glog.V(2).Infof("Got Kube Pod: %+v", pod)
		pod := self.parsePod(&pod)
		addrs, err := net.LookupIP(pod.Hostname)
		if err != nil {
			glog.Errorf("Skipping host %s since looking up its IP failed - %s", pod.Hostname, err)
			self.recordNodeError(pod.Hostname)
			continue
		}
		pod.HostIP = addrs[0].String()
		out = append(out, *pod)
	}

	return out, nil
}

func (self *KubeSource) getStatsFromKubelet(pod Pod, containerName string) (cadvisor.ContainerSpec, []*cadvisor.ContainerStats, error) {
	var containerInfo cadvisor.ContainerInfo
	values := url.Values{}
	values.Add("num_stats", strconv.Itoa(int(time.Since(self.lastQuery)/time.Second)))
	url := "http://" + pod.HostIP + ":" + self.kubeletPort + filepath.Join("/stats", pod.Namespace, pod.Name, pod.ID, containerName) + "?" + values.Encode()
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return cadvisor.ContainerSpec{}, []*cadvisor.ContainerStats{}, err
	}
	err = PostRequestAndGetValue(&http.Client{}, req, &containerInfo)
	if err != nil {
		glog.Errorf("failed to get stats from kubelet url: %s - %s\n", url, err)
		self.recordPodError(pod)
		return cadvisor.ContainerSpec{}, []*cadvisor.ContainerStats{}, nil
	}

	return containerInfo.Spec, containerInfo.Stats, nil
}

func (self *KubeSource) getNodesInfo() ([]RawContainer, error) {
	kubeNodes, err := self.listMinions()
	if err != nil {
		return []RawContainer{}, err
	}
	nodesInfo := []RawContainer{}
	for node, ip := range kubeNodes.Hosts {
		spec, stats, err := self.getStatsFromKubelet(Pod{HostIP: ip}, "/")
		if err != nil {
			glog.V(1).Infof("Failed to get machine stats from kubelet for node %s", node)
			return []RawContainer{}, err
		}
		if len(stats) > 0 {
			container := RawContainer{node, Container{"/", spec, stats}}
			nodesInfo = append(nodesInfo, container)
		}
	}

	return nodesInfo, nil
}

func (self *KubeSource) GetInfo() (ContainerData, error) {
	pods, err := self.getPods()
	if err != nil {
		return ContainerData{}, err
	}
	for _, pod := range pods {
		for _, container := range pod.Containers {
			spec, stats, err := self.getStatsFromKubelet(pod, container.Name)
			if err != nil {
				return ContainerData{}, err
			}
			glog.V(2).Infof("Fetched stats from kubelet for container %s in pod %s", container.Name, pod.Name)
			container.Stats = stats
			container.Spec = spec
		}
	}
	nodesInfo, err := self.getNodesInfo()
	if err != nil {
		return ContainerData{}, err
	}
	glog.V(2).Info("Fetched list of nodes from the master")
	self.lastQuery = time.Now()

	return ContainerData{Pods: pods, Machine: nodesInfo}, nil
}

func newKubeSource() (*KubeSource, error) {
	if len(*argMaster) == 0 {
		return nil, fmt.Errorf("kubernetes_master flag not specified")
	}
	kubeClient := kube_client.NewOrDie(&kube_client.Config{
		Host:     "http://" + *argMaster,
		Version:  kubeClientVersion,
		Insecure: true,
	})

	glog.Infof("Using Kubernetes client with master %q and version %s\n", *argMaster, kubeClientVersion)
	glog.Infof("Using kubelet port %q", *argKubeletPort)

	return &KubeSource{
		client:      kubeClient,
		lastQuery:   time.Now(),
		kubeletPort: *argKubeletPort,
		nodeErrors:  make(map[string]int),
		podErrors:   make(map[PodInstance]int),
	}, nil
}

func (self *KubeSource) GetConfig() string {
	desc := "Source type: Kube\n"
	desc += fmt.Sprintf("\tClient config: master ip %q, version %s\n", *argMaster, kubeClientVersion)
	desc += fmt.Sprintf("\tUsing kubelet port %q\n", self.kubeletPort)
	desc += self.getState()
	desc += "\n"
	return desc
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

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

func TestAccDataSourceAwsNatGateway(t *testing.T) {
	// This is used as a portion of CIDR network addresses.
	rInt := acctest.RandIntRange(4, 254)

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccDataSourceAwsNatGatewayConfig(rInt),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttrPair(
						"data.aws_nat_gateway.test_by_id", "id",
						"aws_nat_gateway.test", "id"),
					resource.TestCheckResourceAttrPair(
						"data.aws_nat_gateway.test_by_subnet_id", "subnet_id",
						"aws_nat_gateway.test", "subnet_id"),
					resource.TestCheckResourceAttrSet("data.aws_nat_gateway.test_by_id", "state"),
					resource.TestCheckNoResourceAttr("data.aws_nat_gateway.test_by_id", "attached_vpc_id"),
				),
			},
		},
	})
}

func testAccDataSourceAwsNatGatewayConfig(rInt int) string {
	return fmt.Sprintf(`
provider "aws" {
  region = "us-west-2"
}

resource "aws_vpc" "test" {
  cidr_block = "172.%d.0.0/16"
  tags {
    Name = "terraform-testacc-nat-gateway-data-source-%d"
  }
}

resource "aws_subnet" "test" {
  vpc_id            = "${aws_vpc.test.id}"
  cidr_block        = "172.%d.123.0/24"
  availability_zone = "us-west-2a"

  tags {
    Name = "terraform-testacc-nat-gateway-data-source-%d"
  }
}

# EIPs are not taggable
resource "aws_eip" "test" {
  vpc = true
}

# IGWs are required for an NGW to spin up; manual dependency
resource "aws_internet_gateway" "test" {
  vpc_id = "${aws_vpc.test.id}"
  tags {
    Name = "terraform-testacc-nat-gateway-data-source-%d"
  }
}

# NGWs are not taggable, either
resource "aws_nat_gateway" "test" {
  subnet_id     = "${aws_subnet.test.id}"
  allocation_id = "${aws_eip.test.id}"

  depends_on = ["aws_internet_gateway.test"]
}

data "aws_nat_gateway" "test_by_id" {
  id = "${aws_nat_gateway.test.id}"
}

data "aws_nat_gateway" "test_by_subnet_id" {
  subnet_id = "${aws_nat_gateway.test.subnet_id}"
}

`, rInt, rInt, rInt, rInt, rInt)
}
<commit_msg>Update acceptance test with new attributes<commit_after>package aws

import (
	"fmt"
	"testing"

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

func TestAccDataSourceAwsNatGateway(t *testing.T) {
	// This is used as a portion of CIDR network addresses.
	rInt := acctest.RandIntRange(4, 254)

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccDataSourceAwsNatGatewayConfig(rInt),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttrPair(
						"data.aws_nat_gateway.test_by_id", "id",
						"aws_nat_gateway.test", "id"),
					resource.TestCheckResourceAttrPair(
						"data.aws_nat_gateway.test_by_subnet_id", "subnet_id",
						"aws_nat_gateway.test", "subnet_id"),
					resource.TestCheckResourceAttrSet("data.aws_nat_gateway.test_by_id", "state"),
					resource.TestCheckResourceAttrSet("data.aws_nat_gateway.test_by_id", "allocation_id"),
					resource.TestCheckResourceAttrSet("data.aws_nat_gateway.test_by_id", "network_interface_id"),
					resource.TestCheckResourceAttrSet("data.aws_nat_gateway.test_by_id", "public_ip"),
					resource.TestCheckResourceAttrSet("data.aws_nat_gateway.test_by_id", "private_ip"),
					resource.TestCheckNoResourceAttr("data.aws_nat_gateway.test_by_id", "attached_vpc_id"),
				),
			},
		},
	})
}

func testAccDataSourceAwsNatGatewayConfig(rInt int) string {
	return fmt.Sprintf(`
provider "aws" {
  region = "us-west-2"
}

resource "aws_vpc" "test" {
  cidr_block = "172.%d.0.0/16"
  tags {
    Name = "terraform-testacc-nat-gateway-data-source-%d"
  }
}

resource "aws_subnet" "test" {
  vpc_id            = "${aws_vpc.test.id}"
  cidr_block        = "172.%d.123.0/24"
  availability_zone = "us-west-2a"

  tags {
    Name = "terraform-testacc-nat-gateway-data-source-%d"
  }
}

# EIPs are not taggable
resource "aws_eip" "test" {
  vpc = true
}

# IGWs are required for an NGW to spin up; manual dependency
resource "aws_internet_gateway" "test" {
  vpc_id = "${aws_vpc.test.id}"
  tags {
    Name = "terraform-testacc-nat-gateway-data-source-%d"
  }
}

# NGWs are not taggable, either
resource "aws_nat_gateway" "test" {
  subnet_id     = "${aws_subnet.test.id}"
  allocation_id = "${aws_eip.test.id}"

  depends_on = ["aws_internet_gateway.test"]
}

data "aws_nat_gateway" "test_by_id" {
  id = "${aws_nat_gateway.test.id}"
}

data "aws_nat_gateway" "test_by_subnet_id" {
  subnet_id = "${aws_nat_gateway.test.subnet_id}"
}

`, rInt, rInt, rInt, rInt, rInt)
}
<|endoftext|>


<commit_before>package shodan

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/assert"
	"time"
)

const (
	testClientToken = "TEST_TOKEN"
	stubsDir        = "stubs"
)

var (
	mux    *http.ServeMux
	server *httptest.Server
	client *Client
)

func setUpTestServe() {
	mux = http.NewServeMux()
	server = httptest.NewServer(mux)
	client = NewClient(nil, testClientToken)
	client.BaseURL = server.URL
	client.ExploitBaseURL = server.URL
	client.StreamBaseURL = server.URL
}

func getStub(t *testing.T, stubName string) []byte {
	stubPath := fmt.Sprintf("%s/%s.json", stubsDir, stubName)
	content, err := ioutil.ReadFile(stubPath)
	if err != nil {
		t.Errorf("getStub error %v", err)
	}

	return content
}

func tearDownTestServe() {
	server.Close()
}

func TestNewClient(t *testing.T) {
	client := NewClient(nil, testClientToken)
	assert.Equal(t, testClientToken, client.Token)
}

func TestNewClient_httpClient(t *testing.T) {
	transport := &http.Transport{Proxy: http.ProxyFromEnvironment}
	httpClient := &http.Client{Transport: transport}
	client := NewClient(httpClient, testClientToken)
	assert.ObjectsAreEqual(httpClient, client.Client)
}

func TestClient_buildURL_success(t *testing.T) {
	client := NewClient(nil, testClientToken)
	testOptions := struct {
		Page    int  `url:"page"`
		ShowAll bool `url:"show_all"`
	}{
		100,
		true,
	}
	testCases := []struct {
		path     string
		params   interface{}
		expected string
	}{
		{
			"/testing/test/1",
			nil,
			baseURL + "/testing/test/1?key=" + testClientToken,
		},
		{
			"/testing/test/2",
			testOptions,
			baseURL + "/testing/test/2?key=" + testClientToken + "&page=100&show_all=true",
		},
	}

	for _, caseParams := range testCases {
		url, err := client.buildURL(baseURL, caseParams.path, caseParams.params)

		assert.Nil(t, err)
		assert.Equal(t, caseParams.expected, url)
	}
}

func TestClient_buildURL_errorBaseURL(t *testing.T) {
	client := NewClient(nil, testClientToken)
	_, err := client.buildURL(":/shodan.io", "", nil)
	assert.NotNil(t, err)
}

func TestClient_buildBaseURL(t *testing.T) {
	expected := client.BaseURL + "/test-base-url-building/?key=" + testClientToken
	actual, err := client.buildBaseURL("/test-base-url-building/", nil)

	assert.Nil(t, err)
	assert.Equal(t, expected, actual)
}

func TestClient_buildExploitBaseURL(t *testing.T) {
	expected := client.ExploitBaseURL + "/test-exploit-url-building/?key=" + testClientToken
	actual, err := client.buildExploitBaseURL("/test-exploit-url-building/", nil)

	assert.Nil(t, err)
	assert.Equal(t, expected, actual)
}

func TestClient_buildStreamBaseURL(t *testing.T) {
	expected := client.BaseURL + "/test-stream-url-building/?key=" + testClientToken
	actual, err := client.buildStreamBaseURL("/test-stream-url-building/", nil)

	assert.Nil(t, err)
	assert.Equal(t, expected, actual)
}

func TestClient_executeRequest_textUnauthorized(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	unauthorizedPath := "/http-error/401"

	errorText := "401 Unauthorized\n\n"
	errorText += "This server could not verify that you are authorized to access the document you requested.  " +
		"Either you supplied the wrong credentials (e.g., bad password), or your browser does not understand how to " +
		"supply the credentials required."

	mux.HandleFunc(unauthorizedPath, func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, errorText, http.StatusUnauthorized)
	})

	url, err := client.buildBaseURL(unauthorizedPath, nil)
	assert.Nil(t, err)

	err = client.executeRequest("GET", url, nil, nil)
	assert.NotNil(t, err)
}

func TestClient_executeRequest_jsonNotFound(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	notFoundPath := "/http-error/404"

	mux.HandleFunc(notFoundPath, func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, `{"error": "No information available for that IP."}`, http.StatusNotFound)
	})

	url, err := client.buildBaseURL(notFoundPath, nil)
	assert.Nil(t, err)

	err = client.executeRequest("GET", url, nil, nil)
	assert.NotNil(t, err)
}

func TestClient_executeStreamRequest_success(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	streamPath := "/stream/success"
	chunkLimit := 3

	mux.HandleFunc(streamPath, func(w http.ResponseWriter, r *http.Request) {
		flusher, ok := w.(http.Flusher)
		if !ok {
			t.Errorf("Cannot use Flush")
		}

		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")

		for i := 0; i < chunkLimit; i++ {
			fmt.Fprintln(w, "chunk")
			flusher.Flush()
			time.Sleep(time.Millisecond * 500)
		}
	})

	url, err := client.buildStreamBaseURL(streamPath, nil)
	assert.Nil(t, err)

	bytesChan := make(chan []byte)
	err = client.executeStreamRequest("GET", url, bytesChan)
	assert.Nil(t, err)

	receivedChunks := 0

	for {
		msg, open := <- bytesChan
		if !open {
			break
		}
		assert.NotEmpty(t, msg)
		receivedChunks++
	}

	assert.Equal(t, chunkLimit, receivedChunks)
}

func TestClient_executeStreamRequest_errorRequest(t *testing.T) {
	url, err := client.buildStreamBaseURL("/stream/error", nil)
	assert.Nil(t, err)

	bytesChan := make(chan []byte)
	err = client.executeStreamRequest("GET", url, bytesChan)
	assert.NotNil(t, err)
}
<commit_msg>small fmt<commit_after>package shodan

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/assert"
	"time"
)

const (
	testClientToken = "TEST_TOKEN"
	stubsDir        = "stubs"
)

var (
	mux    *http.ServeMux
	server *httptest.Server
	client *Client
)

func setUpTestServe() {
	mux = http.NewServeMux()
	server = httptest.NewServer(mux)
	client = NewClient(nil, testClientToken)
	client.BaseURL = server.URL
	client.ExploitBaseURL = server.URL
	client.StreamBaseURL = server.URL
}

func getStub(t *testing.T, stubName string) []byte {
	stubPath := fmt.Sprintf("%s/%s.json", stubsDir, stubName)
	content, err := ioutil.ReadFile(stubPath)
	if err != nil {
		t.Errorf("getStub error %v", err)
	}

	return content
}

func tearDownTestServe() {
	server.Close()
}

func TestNewClient(t *testing.T) {
	client := NewClient(nil, testClientToken)
	assert.Equal(t, testClientToken, client.Token)
}

func TestNewClient_httpClient(t *testing.T) {
	transport := &http.Transport{Proxy: http.ProxyFromEnvironment}
	httpClient := &http.Client{Transport: transport}
	client := NewClient(httpClient, testClientToken)
	assert.ObjectsAreEqual(httpClient, client.Client)
}

func TestClient_buildURL_success(t *testing.T) {
	client := NewClient(nil, testClientToken)
	testOptions := struct {
		Page    int  `url:"page"`
		ShowAll bool `url:"show_all"`
	}{
		100,
		true,
	}
	testCases := []struct {
		path     string
		params   interface{}
		expected string
	}{
		{
			"/testing/test/1",
			nil,
			baseURL + "/testing/test/1?key=" + testClientToken,
		},
		{
			"/testing/test/2",
			testOptions,
			baseURL + "/testing/test/2?key=" + testClientToken + "&page=100&show_all=true",
		},
	}

	for _, caseParams := range testCases {
		url, err := client.buildURL(baseURL, caseParams.path, caseParams.params)

		assert.Nil(t, err)
		assert.Equal(t, caseParams.expected, url)
	}
}

func TestClient_buildURL_errorBaseURL(t *testing.T) {
	client := NewClient(nil, testClientToken)
	_, err := client.buildURL(":/shodan.io", "", nil)
	assert.NotNil(t, err)
}

func TestClient_buildBaseURL(t *testing.T) {
	expected := client.BaseURL + "/test-base-url-building/?key=" + testClientToken
	actual, err := client.buildBaseURL("/test-base-url-building/", nil)

	assert.Nil(t, err)
	assert.Equal(t, expected, actual)
}

func TestClient_buildExploitBaseURL(t *testing.T) {
	expected := client.ExploitBaseURL + "/test-exploit-url-building/?key=" + testClientToken
	actual, err := client.buildExploitBaseURL("/test-exploit-url-building/", nil)

	assert.Nil(t, err)
	assert.Equal(t, expected, actual)
}

func TestClient_buildStreamBaseURL(t *testing.T) {
	expected := client.BaseURL + "/test-stream-url-building/?key=" + testClientToken
	actual, err := client.buildStreamBaseURL("/test-stream-url-building/", nil)

	assert.Nil(t, err)
	assert.Equal(t, expected, actual)
}

func TestClient_executeRequest_textUnauthorized(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	unauthorizedPath := "/http-error/401"

	errorText := "401 Unauthorized\n\n"
	errorText += "This server could not verify that you are authorized to access the document you requested.  " +
		"Either you supplied the wrong credentials (e.g., bad password), or your browser does not understand how to " +
		"supply the credentials required."

	mux.HandleFunc(unauthorizedPath, func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, errorText, http.StatusUnauthorized)
	})

	url, err := client.buildBaseURL(unauthorizedPath, nil)
	assert.Nil(t, err)

	err = client.executeRequest("GET", url, nil, nil)
	assert.NotNil(t, err)
}

func TestClient_executeRequest_jsonNotFound(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	notFoundPath := "/http-error/404"

	mux.HandleFunc(notFoundPath, func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, `{"error": "No information available for that IP."}`, http.StatusNotFound)
	})

	url, err := client.buildBaseURL(notFoundPath, nil)
	assert.Nil(t, err)

	err = client.executeRequest("GET", url, nil, nil)
	assert.NotNil(t, err)
}

func TestClient_executeStreamRequest_success(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	streamPath := "/stream/success"
	chunkLimit := 3

	mux.HandleFunc(streamPath, func(w http.ResponseWriter, r *http.Request) {
		flusher, ok := w.(http.Flusher)
		if !ok {
			t.Errorf("Cannot use Flush")
		}

		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")

		for i := 0; i < chunkLimit; i++ {
			fmt.Fprintln(w, "chunk")
			flusher.Flush()
			time.Sleep(time.Millisecond * 500)
		}
	})

	url, err := client.buildStreamBaseURL(streamPath, nil)
	assert.Nil(t, err)

	bytesChan := make(chan []byte)
	err = client.executeStreamRequest("GET", url, bytesChan)
	assert.Nil(t, err)

	receivedChunks := 0

	for {
		msg, open := <-bytesChan
		if !open {
			break
		}
		assert.NotEmpty(t, msg)
		receivedChunks++
	}

	assert.Equal(t, chunkLimit, receivedChunks)
}

func TestClient_executeStreamRequest_errorRequest(t *testing.T) {
	url, err := client.buildStreamBaseURL("/stream/error", nil)
	assert.Nil(t, err)

	bytesChan := make(chan []byte)
	err = client.executeStreamRequest("GET", url, bytesChan)
	assert.NotNil(t, err)
}
<|endoftext|>


<commit_before>package azurerm

import (
	"fmt"
	"strings"

	"github.com/Azure/azure-sdk-for-go/arm/containerinstance"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceArmContainerGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceArmContainerGroupCreate,
		Read:   resourceArmContainerGroupRead,
		Update: resourceArmContainerGroupCreate,
		Delete: resourceArmContainerGroupDelete,

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"location": locationSchema(),

			"resource_group_name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"image": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"cpu": {
				Type:     schema.TypeFloat,
				Optional: true,
				ForceNew: true,
			},

			"ip_address_type": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"memory": {
				Type:     schema.TypeFloat,
				Optional: true,
				ForceNew: true,
			},

			"os_type": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"port": {
				Type:     schema.TypeInt,
				Optional: true,
				ForceNew: true,
			},

			"protocol": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"ip_address": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceArmContainerGroupCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient)
	containerGroupsClient := client.containerGroupsClient

	// container group properties
	resGroup := d.Get("resource_group_name").(string)
	name := d.Get("name").(string)
	location := d.Get("location").(string)
	OSType := d.Get("os_type").(string)
	IPAddressType := d.Get("ip_address_type").(string)
	protocol := d.Get("protocol").(string)
	tags := d.Get("tags").(map[string]interface{})

	// per container properties
	image := d.Get("image").(string)
	cpu := d.Get("cpu").(float64)
	memory := d.Get("memory").(float64)
	port := int32(d.Get("port").(int))

	// type ContainerGroupProperties struct {
	// 	ProvisioningState        *string                    `json:"provisioningState,omitempty"`
	// 	Containers               *[]Container               `json:"containers,omitempty"`
	// 	ImageRegistryCredentials *[]ImageRegistryCredential `json:"imageRegistryCredentials,omitempty"`
	// 	RestartPolicy            ContainerRestartPolicy     `json:"restartPolicy,omitempty"`
	// 	IPAddress                *IPAddress                 `json:"ipAddress,omitempty"`
	// 	OsType                   OperatingSystemTypes       `json:"osType,omitempty"`
	// 	State                    *string                    `json:"state,omitempty"`
	// 	Volumes                  *[]Volume                  `json:"volumes,omitempty"`
	// }

	// type ContainerProperties struct {
	// 	Image                *string                          `json:"image,omitempty"`
	// 	Command              *[]string                        `json:"command,omitempty"`
	// 	Ports                *[]ContainerPort                 `json:"ports,omitempty"`
	// 	EnvironmentVariables *[]EnvironmentVariable           `json:"environmentVariables,omitempty"`
	// 	InstanceView         *ContainerPropertiesInstanceView `json:"instanceView,omitempty"`
	// 	Resources            *ResourceRequirements            `json:"resources,omitempty"`
	// 	VolumeMounts         *[]VolumeMount                   `json:"volumeMounts,omitempty"`
	// }

	// per container port (port number only)
	containerPort := containerinstance.ContainerPort{
		Port: &port,
	}

	container := containerinstance.Container{
		Name: &name,
		ContainerProperties: &containerinstance.ContainerProperties{
			Image: &image,
			Ports: &[]containerinstance.ContainerPort{containerPort},
			Resources: &containerinstance.ResourceRequirements{
				Requests: &containerinstance.ResourceRequests{
					MemoryInGB: &memory,
					CPU:        &cpu,
				},
			},
		},
	}

	// container group port (port number + protocol)
	containerGroupPort := containerinstance.Port{
		Port: &port,
	}

	if strings.ToUpper(protocol) == "TCP" || strings.ToUpper(protocol) == "UDP" {
		containerGroupPort.Protocol = containerinstance.ContainerGroupNetworkProtocol(strings.ToUpper(protocol))
	}

	containerGroup := containerinstance.ContainerGroup{
		Name:     &name,
		Location: &location,
		Tags:     expandTags(tags),
		ContainerGroupProperties: &containerinstance.ContainerGroupProperties{
			Containers: &[]containerinstance.Container{container},
			IPAddress: &containerinstance.IPAddress{
				Type:  &IPAddressType,
				Ports: &[]containerinstance.Port{containerGroupPort},
			},
			OsType: containerinstance.OperatingSystemTypes(OSType),
		},
	}

	_, error := containerGroupsClient.CreateOrUpdate(resGroup, name, containerGroup)
	if error != nil {
		return error
	}

	read, readErr := containerGroupsClient.Get(resGroup, name)
	if readErr != nil {
		return readErr
	}

	if read.ID == nil {
		return fmt.Errorf("Cannot read container group %s (resource group %s) ID", name, resGroup)
	}

	d.SetId(*read.ID)

	return nil
}
func resourceArmContainerGroupRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient)
	containterGroupsClient := client.containerGroupsClient

	id, err := parseAzureResourceID(d.Id())

	if err != nil {
		return err
	}

	resGroup := id.ResourceGroup
	name := id.Path["containerGroups"]

	resp, error := containterGroupsClient.Get(resGroup, name)

	if error != nil {
		return error
	}

	d.Set("name", name)
	d.Set("resource_group_name", resGroup)
	d.Set("location", azureRMNormalizeLocation(*resp.Location))
	flattenAndSetTags(d, resp.Tags)

	d.Set("os_type", string(resp.OsType))
	d.Set("ip_address_type", *resp.IPAddress.Type)
	d.Set("ip_address", *resp.IPAddress.IP)

	ports := *resp.IPAddress.Ports
	d.Set("protocol", string(ports[0].Protocol))

	containers := *resp.Containers
	d.Set("image", containers[0].Image)
	resourceRequirements := *containers[0].Resources
	resourceRequests := *resourceRequirements.Requests
	d.Set("cpu", *resourceRequests.CPU)
	d.Set("memory", *resourceRequests.MemoryInGB)
	containerPorts := *containers[0].Ports
	d.Set("port", containerPorts[0].Port)

	return nil
}
func resourceArmContainerGroupDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient)
	containterGroupsClient := client.containerGroupsClient

	// container group properties
	resGroup := d.Get("resource_group_name").(string)
	name := d.Get("name").(string)

	_, error := containterGroupsClient.Delete(resGroup, name)

	if error != nil {
		return error
	}

	return nil
}
<commit_msg>WIP Container Groups resource<commit_after>package azurerm

import (
	"fmt"
	"strings"

	"github.com/Azure/azure-sdk-for-go/arm/containerinstance"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceArmContainerGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceArmContainerGroupCreate,
		Read:   resourceArmContainerGroupRead,
		Update: resourceArmContainerGroupCreate,
		Delete: resourceArmContainerGroupDelete,

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"location": locationSchema(),

			"resource_group_name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"ip_address_type": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"os_type": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"tags": tagsSchema(),

			"ip_address": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"container": {
				Type:     schema.TypeList,
				Required: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{

						"name": {
							Type:     schema.TypeString,
							Required: true,
							ForceNew: true,
						},

						"image": {
							Type:     schema.TypeString,
							Required: true,
							ForceNew: true,
						},

						"cpu": {
							Type:     schema.TypeFloat,
							Required: true,
							ForceNew: true,
						},

						"memory": {
							Type:     schema.TypeFloat,
							Required: true,
							ForceNew: true,
						},

						"port": {
							Type:     schema.TypeInt,
							Required: true,
							ForceNew: true,
						},

						"protocol": {
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: true,
						},
					},
				},
			},
		},
	}
}

func resourceArmContainerGroupCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient)
	containerGroupsClient := client.containerGroupsClient

	// container group properties
	resGroup := d.Get("resource_group_name").(string)
	name := d.Get("name").(string)
	location := d.Get("location").(string)
	OSType := d.Get("os_type").(string)
	if OSType == "" {
		OSType = "linux"
	}
	IPAddressType := d.Get("ip_address_type").(string)
	if IPAddressType == "" {
		IPAddressType = "public"
	}
	tags := d.Get("tags").(map[string]interface{})

	containersConfig := d.Get("container").([]interface{})
	containers := make([]containerinstance.Container, len(containersConfig))
	containerGroupPorts := make([]containerinstance.Port, len(containersConfig))

	for index, containerConfig := range containersConfig {
		data := containerConfig.(map[string]interface{})

		// required
		name := data["name"].(string)
		image := data["image"].(string)

		// optional
		cpu := data["cpu"].(float64)
		memory := data["memory"].(float64)
		port := int32(data["port"].(int))
		protocol := data["protocol"].(string)

		containerPort := containerinstance.ContainerPort{
			Port: &port,
		}

		// container group port (port number + protocol)
		containerGroupPort := containerinstance.Port{
			Port: &port,
		}

		if strings.ToUpper(protocol) == "TCP" || strings.ToUpper(protocol) == "UDP" {
			containerGroupPort.Protocol = containerinstance.ContainerGroupNetworkProtocol(strings.ToUpper(protocol))
		}
		containerGroupPorts[index] = containerGroupPort

		container := containerinstance.Container{
			Name: &name,
			ContainerProperties: &containerinstance.ContainerProperties{
				Image: &image,
				Ports: &[]containerinstance.ContainerPort{containerPort},
				Resources: &containerinstance.ResourceRequirements{
					Requests: &containerinstance.ResourceRequests{
						MemoryInGB: &memory,
						CPU:        &cpu,
					},
				},
			},
		}

		containers[index] = container
	}

	// type ContainerGroupProperties struct {
	// 	ProvisioningState        *string                    `json:"provisioningState,omitempty"`
	// 	Containers               *[]Container               `json:"containers,omitempty"`
	// 	ImageRegistryCredentials *[]ImageRegistryCredential `json:"imageRegistryCredentials,omitempty"`
	// 	RestartPolicy            ContainerRestartPolicy     `json:"restartPolicy,omitempty"`
	// 	IPAddress                *IPAddress                 `json:"ipAddress,omitempty"`
	// 	OsType                   OperatingSystemTypes       `json:"osType,omitempty"`
	// 	State                    *string                    `json:"state,omitempty"`
	// 	Volumes                  *[]Volume                  `json:"volumes,omitempty"`
	// }

	// type ContainerProperties struct {
	// 	Image                *string                          `json:"image,omitempty"`
	// 	Command              *[]string                        `json:"command,omitempty"`
	// 	Ports                *[]ContainerPort                 `json:"ports,omitempty"`
	// 	EnvironmentVariables *[]EnvironmentVariable           `json:"environmentVariables,omitempty"`
	// 	InstanceView         *ContainerPropertiesInstanceView `json:"instanceView,omitempty"`
	// 	Resources            *ResourceRequirements            `json:"resources,omitempty"`
	// 	VolumeMounts         *[]VolumeMount                   `json:"volumeMounts,omitempty"`
	// }

	containerGroup := containerinstance.ContainerGroup{
		Name:     &name,
		Location: &location,
		Tags:     expandTags(tags),
		ContainerGroupProperties: &containerinstance.ContainerGroupProperties{
			Containers: &containers,
			IPAddress: &containerinstance.IPAddress{
				Type:  &IPAddressType,
				Ports: &containerGroupPorts,
			},
			OsType: containerinstance.OperatingSystemTypes(OSType),
		},
	}

	_, error := containerGroupsClient.CreateOrUpdate(resGroup, name, containerGroup)
	if error != nil {
		return error
	}

	read, readErr := containerGroupsClient.Get(resGroup, name)
	if readErr != nil {
		return readErr
	}

	if read.ID == nil {
		return fmt.Errorf("Cannot read container group %s (resource group %s) ID", name, resGroup)
	}

	d.SetId(*read.ID)

	return nil
}
func resourceArmContainerGroupRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient)
	containterGroupsClient := client.containerGroupsClient

	id, err := parseAzureResourceID(d.Id())

	if err != nil {
		return err
	}

	resGroup := id.ResourceGroup
	name := id.Path["containerGroups"]

	resp, error := containterGroupsClient.Get(resGroup, name)

	if error != nil {
		return error
	}

	d.Set("name", name)
	d.Set("resource_group_name", resGroup)
	d.Set("location", azureRMNormalizeLocation(*resp.Location))
	flattenAndSetTags(d, resp.Tags)

	d.Set("os_type", string(resp.OsType))
	d.Set("ip_address_type", *resp.IPAddress.Type)
	d.Set("ip_address", *resp.IPAddress.IP)

	// ports := *resp.IPAddress.Ports
	// containers := *resp.Containers

	return nil
}
func resourceArmContainerGroupDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient)
	containterGroupsClient := client.containerGroupsClient

	// container group properties
	resGroup := d.Get("resource_group_name").(string)
	name := d.Get("name").(string)

	_, error := containterGroupsClient.Delete(resGroup, name)

	if error != nil {
		return error
	}

	return nil
}
<|endoftext|>


<commit_before>package yang

import (
	"github.com/openconfig/goyang/pkg/yang"
	"github.com/oshothebig/pbast"
)

var builtinMap = map[yang.TypeKind]pbast.Type{
	yang.Yint8:   pbast.Int32,
	yang.Yint16:  pbast.Int32,
	yang.Yint32:  pbast.Int32,
	yang.Yint64:  pbast.Int64,
	yang.Yuint8:  pbast.UInt32,
	yang.Yuint16: pbast.UInt32,
	yang.Yuint32: pbast.UInt32,
	yang.Yuint64: pbast.UInt64,
	yang.Ystring: pbast.String,
	yang.Ybool:   pbast.Bool,
	yang.Ybinary: pbast.Bytes,
}

type transformer struct {
	topScope  []*pbast.Message
	decimal64 *pbast.Message
}

// e must be YANG module
func Transform(e *yang.Entry) *pbast.File {
	if _, ok := e.Node.(*yang.Module); !ok {
		return nil
	}

	t := &transformer{}

	return t.module(entry{e})
}

func (t *transformer) declare(m *pbast.Message) {
	if m == nil {
		return
	}
	t.topScope = append(t.topScope, m)
}

func (t *transformer) module(e entry) *pbast.File {
	namespace := e.Namespace().Name
	f := pbast.NewFile(pbast.NewPackageWithElements(guessElements(namespace)))

	f.Comment = t.moduleComment(e)

	// Child nodes are enclosed with Root message
	f.AddMessage(t.buildMessage("Root", e))

	// RPCs
	s := t.rpcs(e)
	f.AddService(s)

	// Notifications
	n := t.notifications(e)
	f.AddService(n)

	for _, m := range t.topScope {
		f.AddMessage(m)
	}
	f.AddMessage(t.decimal64)

	return f
}

func (t *transformer) moduleComment(e entry) pbast.Comment {
	description := t.description(e)
	namespace := t.namespace(e)
	revisions := t.revisions(e)
	reference := t.reference(e)

	var comment []string
	comment = append(comment, description...)
	comment = append(comment, namespace...)
	comment = append(comment, revisions...)
	comment = append(comment, reference...)

	return comment
}

func (t *transformer) genericComments(e entry) pbast.Comment {
	description := t.description(e)
	reference := t.reference(e)

	comments := append(description, reference...)
	return comments
}

func (t *transformer) description(e entry) pbast.Comment {
	description := e.Description
	if e.Description == "" {
		return nil
	}

	return []string{"Reference:", description}
}

func (t *transformer) revisions(e entry) pbast.Comment {
	var lines []string
	if v := e.Extra["revision"]; len(v) > 0 {
		for _, rev := range v[0].([]*yang.Revision) {
			lines = append(lines, "Revision: "+rev.Name)
		}
	}

	return lines
}

func (t *transformer) namespace(e entry) pbast.Comment {
	namespace := e.Namespace().Name
	if namespace == "" {
		return nil
	}

	return []string{"Namespace: " + namespace}
}

func (t *transformer) reference(e entry) pbast.Comment {
	v := e.Extra["reference"]
	if len(v) == 0 {
		return nil
	}

	ref := v[0].(*yang.Value)
	if ref == nil {
		return nil
	}
	if ref.Name == "" {
		return nil
	}

	return []string{"Reference:", ref.Name}
}

func (t *transformer) rpcs(e entry) *pbast.Service {
	rpcs := e.rpcs()
	if len(rpcs) == 0 {
		return nil
	}

	s := pbast.NewService(CamelCase(e.Name))
	s.Comment = t.genericComments(e)
	for _, rpc := range rpcs {
		r := t.rpc(rpc)
		s.AddRPC(r)
	}

	return s
}

func (t *transformer) rpc(e entry) *pbast.RPC {
	method := CamelCase(e.Name)
	in := method + "Request"
	out := method + "Response"

	rpc := pbast.NewRPC(
		method,
		pbast.NewReturnType(in),
		pbast.NewReturnType(out),
	)
	rpc.Comment = t.genericComments(e)

	t.declare(t.buildMessage(in, entry{e.RPC.Input}))
	t.declare(t.buildMessage(out, entry{e.RPC.Output}))

	return rpc
}

func (t *transformer) notifications(e entry) *pbast.Service {
	notifications := e.notifications()
	if len(notifications) == 0 {
		return nil
	}

	s := pbast.NewService(CamelCase(e.Name + "Notification"))
	s.Comment = t.genericComments(e)
	for _, notification := range notifications {
		n := t.notification(notification)
		s.AddRPC(n)
	}

	return s
}

func (t *transformer) notification(e entry) *pbast.RPC {
	method := CamelCase(e.Name)
	in := method + "NotificationRequest"
	out := method + "NotificationResponse"

	rpc := pbast.NewRPC(method, pbast.NewReturnType(in), pbast.NewReturnType(out))

	// notification statement doesn't have an input parameter equivalent,
	// then empty message is used for input as RPC
	t.declare(pbast.NewMessage(in))
	t.declare(t.buildMessage(out, e))

	return rpc
}

func (t *transformer) buildMessage(name string, e entry) *pbast.Message {
	if e.Entry == nil {
		return nil
	}

	msg := pbast.NewMessage(name)
	msg.Comment = t.genericComments(e)
	for index, child := range e.children() {
		fieldNum := index + 1
		switch {
		// leaf-list case
		case child.Type != nil && child.ListAttr != nil:
			field, nested := t.leaf(child, fieldNum, true)
			msg.AddType(nested).AddField(field)
		// leaf case
		case child.Type != nil:
			field, nested := t.leaf(child, fieldNum, false)
			msg.AddType(nested).AddField(field)
		// list case
		case child.ListAttr != nil:
			inner, field := t.directory(child, fieldNum, true)
			msg.AddMessage(inner).AddField(field)
		// others might be container case
		default:
			inner, field := t.directory(child, fieldNum, false)
			msg.AddMessage(inner).AddField(field)
		}
	}

	return msg
}

func (t *transformer) leaf(e entry, index int, repeated bool) (field *pbast.MessageField, nested pbast.Type) {
	typ := builtinMap[e.Type.Kind]
	// no direct builtin type mapping
	// custom message is built
	if typ == nil {
		name := CamelCase(e.Name)
		switch e.Type.Kind {
		// define at the top level
		case yang.Ydecimal64:
			t.decimal64 = decimal64Message
			typ = decimal64Message
		// define as a nested type
		case yang.Ybits:
			typ = t.customBits(name, e.Type.Bit)
		// define as a nested type
		case yang.Yenum:
			typ = t.customEnum(name, e.Type.Enum)
		// not implemented
		case yang.Yunion, yang.Yempty, yang.Yleafref,
			yang.Yidentityref, yang.YinstanceIdentifier:
			return nil, nil
		}
	}

	name := underscoreCase(e.Name)
	if repeated {
		field = pbast.NewRepeatedMessageField(typ, name, index)
	} else {
		field = pbast.NewMessageField(typ, name, index)
	}

	if e.Type.Kind == yang.Ydecimal64 {
		return field, nil
	}

	return field, typ
}

func (t *transformer) customBits(name string, bits *yang.EnumType) *pbast.Message {
	msg := pbast.NewMessage(name)
	for i, n := range bits.Names() {
		v := 1 << uint(bits.Values()[i])
		msg.AddField(pbast.NewMessageField(pbast.Bool, n, v))
	}

	return msg
}

func (t *transformer) customEnum(name string, e *yang.EnumType) *pbast.Enum {
	enum := pbast.NewEnum(name)
	for i, n := range e.Names() {
		v := int(e.Values()[i])
		enum.AddField(pbast.NewEnumField(constantName(n), v))
	}

	return enum
}

func (t *transformer) directory(e entry, index int, repeated bool) (*pbast.Message, *pbast.MessageField) {
	fieldName := underscoreCase(e.Name)
	typeName := CamelCase(e.Name)

	inner := t.buildMessage(typeName, e)
	var field *pbast.MessageField
	if repeated {
		field = pbast.NewRepeatedMessageField(inner, fieldName, index)
	} else {
		field = pbast.NewMessageField(inner, fieldName, index)
	}

	return inner, field
}
<commit_msg>Fix incorrect line break in comments<commit_after>package yang

import (
	"strings"

	"github.com/openconfig/goyang/pkg/yang"
	"github.com/oshothebig/pbast"
)

var builtinMap = map[yang.TypeKind]pbast.Type{
	yang.Yint8:   pbast.Int32,
	yang.Yint16:  pbast.Int32,
	yang.Yint32:  pbast.Int32,
	yang.Yint64:  pbast.Int64,
	yang.Yuint8:  pbast.UInt32,
	yang.Yuint16: pbast.UInt32,
	yang.Yuint32: pbast.UInt32,
	yang.Yuint64: pbast.UInt64,
	yang.Ystring: pbast.String,
	yang.Ybool:   pbast.Bool,
	yang.Ybinary: pbast.Bytes,
}

type transformer struct {
	topScope  []*pbast.Message
	decimal64 *pbast.Message
}

// e must be YANG module
func Transform(e *yang.Entry) *pbast.File {
	if _, ok := e.Node.(*yang.Module); !ok {
		return nil
	}

	t := &transformer{}

	return t.module(entry{e})
}

func (t *transformer) declare(m *pbast.Message) {
	if m == nil {
		return
	}
	t.topScope = append(t.topScope, m)
}

func (t *transformer) module(e entry) *pbast.File {
	namespace := e.Namespace().Name
	f := pbast.NewFile(pbast.NewPackageWithElements(guessElements(namespace)))

	f.Comment = t.moduleComment(e)

	// Child nodes are enclosed with Root message
	f.AddMessage(t.buildMessage("Root", e))

	// RPCs
	s := t.rpcs(e)
	f.AddService(s)

	// Notifications
	n := t.notifications(e)
	f.AddService(n)

	for _, m := range t.topScope {
		f.AddMessage(m)
	}
	f.AddMessage(t.decimal64)

	return f
}

func (t *transformer) moduleComment(e entry) pbast.Comment {
	description := t.description(e)
	namespace := t.namespace(e)
	revisions := t.revisions(e)
	reference := t.reference(e)

	var comment []string
	comment = append(comment, description...)
	comment = append(comment, namespace...)
	comment = append(comment, revisions...)
	comment = append(comment, reference...)

	return comment
}

func (t *transformer) genericComments(e entry) pbast.Comment {
	description := t.description(e)
	reference := t.reference(e)

	comments := append(description, reference...)
	return comments
}

func (t *transformer) description(e entry) pbast.Comment {
	description := e.Description
	if e.Description == "" {
		return nil
	}

	lines := strings.Split(strings.TrimRight(description, "\n "), "\n")

	ret := make([]string, 0, len(lines)+1)
	ret = append(ret, "Description:")
	ret = append(ret, lines...)
	return ret
}

func (t *transformer) revisions(e entry) pbast.Comment {
	var lines []string
	if v := e.Extra["revision"]; len(v) > 0 {
		for _, rev := range v[0].([]*yang.Revision) {
			lines = append(lines, "Revision: "+rev.Name)
		}
	}

	return lines
}

func (t *transformer) namespace(e entry) pbast.Comment {
	namespace := e.Namespace().Name
	if namespace == "" {
		return nil
	}

	return []string{"Namespace: " + namespace}
}

func (t *transformer) reference(e entry) pbast.Comment {
	v := e.Extra["reference"]
	if len(v) == 0 {
		return nil
	}

	ref := v[0].(*yang.Value)
	if ref == nil {
		return nil
	}
	if ref.Name == "" {
		return nil
	}

	lines := strings.Split(strings.TrimRight(ref.Name, "\n "), "\n")

	ret := make([]string, 0, len(lines)+1)
	ret = append(ret, "Reference:")
	ret = append(ret, lines...)
	return ret
}

func (t *transformer) rpcs(e entry) *pbast.Service {
	rpcs := e.rpcs()
	if len(rpcs) == 0 {
		return nil
	}

	s := pbast.NewService(CamelCase(e.Name))
	s.Comment = t.genericComments(e)
	for _, rpc := range rpcs {
		r := t.rpc(rpc)
		s.AddRPC(r)
	}

	return s
}

func (t *transformer) rpc(e entry) *pbast.RPC {
	method := CamelCase(e.Name)
	in := method + "Request"
	out := method + "Response"

	rpc := pbast.NewRPC(
		method,
		pbast.NewReturnType(in),
		pbast.NewReturnType(out),
	)
	rpc.Comment = t.genericComments(e)

	t.declare(t.buildMessage(in, entry{e.RPC.Input}))
	t.declare(t.buildMessage(out, entry{e.RPC.Output}))

	return rpc
}

func (t *transformer) notifications(e entry) *pbast.Service {
	notifications := e.notifications()
	if len(notifications) == 0 {
		return nil
	}

	s := pbast.NewService(CamelCase(e.Name + "Notification"))
	s.Comment = t.genericComments(e)
	for _, notification := range notifications {
		n := t.notification(notification)
		s.AddRPC(n)
	}

	return s
}

func (t *transformer) notification(e entry) *pbast.RPC {
	method := CamelCase(e.Name)
	in := method + "NotificationRequest"
	out := method + "NotificationResponse"

	rpc := pbast.NewRPC(method, pbast.NewReturnType(in), pbast.NewReturnType(out))

	// notification statement doesn't have an input parameter equivalent,
	// then empty message is used for input as RPC
	t.declare(pbast.NewMessage(in))
	t.declare(t.buildMessage(out, e))

	return rpc
}

func (t *transformer) buildMessage(name string, e entry) *pbast.Message {
	if e.Entry == nil {
		return nil
	}

	msg := pbast.NewMessage(name)
	msg.Comment = t.genericComments(e)
	for index, child := range e.children() {
		fieldNum := index + 1
		switch {
		// leaf-list case
		case child.Type != nil && child.ListAttr != nil:
			field, nested := t.leaf(child, fieldNum, true)
			msg.AddType(nested).AddField(field)
		// leaf case
		case child.Type != nil:
			field, nested := t.leaf(child, fieldNum, false)
			msg.AddType(nested).AddField(field)
		// list case
		case child.ListAttr != nil:
			inner, field := t.directory(child, fieldNum, true)
			msg.AddMessage(inner).AddField(field)
		// others might be container case
		default:
			inner, field := t.directory(child, fieldNum, false)
			msg.AddMessage(inner).AddField(field)
		}
	}

	return msg
}

func (t *transformer) leaf(e entry, index int, repeated bool) (field *pbast.MessageField, nested pbast.Type) {
	typ := builtinMap[e.Type.Kind]
	// no direct builtin type mapping
	// custom message is built
	if typ == nil {
		name := CamelCase(e.Name)
		switch e.Type.Kind {
		// define at the top level
		case yang.Ydecimal64:
			t.decimal64 = decimal64Message
			typ = decimal64Message
		// define as a nested type
		case yang.Ybits:
			typ = t.customBits(name, e.Type.Bit)
		// define as a nested type
		case yang.Yenum:
			typ = t.customEnum(name, e.Type.Enum)
		// not implemented
		case yang.Yunion, yang.Yempty, yang.Yleafref,
			yang.Yidentityref, yang.YinstanceIdentifier:
			return nil, nil
		}
	}

	name := underscoreCase(e.Name)
	if repeated {
		field = pbast.NewRepeatedMessageField(typ, name, index)
	} else {
		field = pbast.NewMessageField(typ, name, index)
	}

	if e.Type.Kind == yang.Ydecimal64 {
		return field, nil
	}

	return field, typ
}

func (t *transformer) customBits(name string, bits *yang.EnumType) *pbast.Message {
	msg := pbast.NewMessage(name)
	for i, n := range bits.Names() {
		v := 1 << uint(bits.Values()[i])
		msg.AddField(pbast.NewMessageField(pbast.Bool, n, v))
	}

	return msg
}

func (t *transformer) customEnum(name string, e *yang.EnumType) *pbast.Enum {
	enum := pbast.NewEnum(name)
	for i, n := range e.Names() {
		v := int(e.Values()[i])
		enum.AddField(pbast.NewEnumField(constantName(n), v))
	}

	return enum
}

func (t *transformer) directory(e entry, index int, repeated bool) (*pbast.Message, *pbast.MessageField) {
	fieldName := underscoreCase(e.Name)
	typeName := CamelCase(e.Name)

	inner := t.buildMessage(typeName, e)
	var field *pbast.MessageField
	if repeated {
		field = pbast.NewRepeatedMessageField(inner, fieldName, index)
	} else {
		field = pbast.NewMessageField(inner, fieldName, index)
	}

	return inner, field
}
<|endoftext|>


<commit_before>// Copyright 2016 IBM Corporation
//
//   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 dns

import (
	"fmt"
	"net"
	"net/url"
	"strconv"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/amalgam8/amalgam8/registry/client"
	"github.com/miekg/dns"
)

// Server represent a DNS server. has config field for port,domain,and client discovery, and the DNS server itself
type Server struct {
	config    Config
	dnsServer *dns.Server
}

// Config represents the DNS server configurations.
type Config struct {
	DiscoveryClient client.Discovery
	Port            uint16
	Domain          string
}

// NewServer creates a new instance of a DNS server with the given configurations
func NewServer(config Config) (*Server, error) {
	err := validate(&config)
	if err != nil {
		return nil, err
	}
	s := &Server{
		config: config,
	}

	// Setup DNS muxing
	mux := dns.NewServeMux()
	mux.HandleFunc(config.Domain, s.handleRequest)

	// Setup a DNS server
	s.dnsServer = &dns.Server{
		Addr:    fmt.Sprintf(":%d", config.Port),
		Net:     "udp",
		Handler: mux,
	}

	return s, nil
}

// ListenAndServe starts the DNS server
func (s *Server) ListenAndServe() error {
	logrus.Info("Starting DNS server")
	err := s.dnsServer.ListenAndServe()

	if err != nil {
		logrus.WithError(err).Errorf("Error starting DNS server")
	}

	return nil
}

// Shutdown stops the DNS server
func (s *Server) Shutdown() error {
	logrus.Info("Shutting down DNS server")
	err := s.dnsServer.Shutdown()

	if err != nil {
		logrus.WithError(err).Errorf("Error shutting down DNS server")
	} else {
		logrus.Info("DNS server has shutdown")
	}

	return err
}

func (s *Server) handleRequest(w dns.ResponseWriter, request *dns.Msg) {
	response := new(dns.Msg)
	response.SetReply(request)
	response.Extra = request.Extra
	response.Authoritative = true
	response.RecursionAvailable = false

	for i, question := range request.Question {
		err := s.handleQuestion(question, request, response)
		if err != nil {
			logrus.WithError(err).Errorf("Error handling DNS question %d: %s", i, question.String())
			// TODO: what should the dns response return ?
			break
		}
	}
	err := w.WriteMsg(response)
	if err != nil {
		logrus.WithError(err).Errorf("Error writing DNS response")
	}
}

func (s *Server) handleQuestion(question dns.Question, request, response *dns.Msg) error {

	switch question.Qclass {
	case dns.ClassINET:
	default:
		response.SetRcode(request, dns.RcodeServerFailure)
		return fmt.Errorf("unsupported DNS question class: %v", dns.Class(question.Qclass).String())
	}

	switch question.Qtype {
	case dns.TypeA:
	case dns.TypeAAAA:
	case dns.TypeSRV:
	default:
		response.SetRcode(request, dns.RcodeServerFailure)
		return fmt.Errorf("unsupported DNS question type: %v", dns.Type(question.Qtype).String())
	}

	serviceInstances, err := s.retrieveServices(question, request, response)

	if err != nil {
		return err
	}
	err = s.createRecordsForInstances(question, request, response, serviceInstances)
	return err

}

func (s *Server) retrieveServices(question dns.Question, request, response *dns.Msg) ([]*client.ServiceInstance, error) {
	var serviceInstances []*client.ServiceInstance
	var err error
	// parse query :
	// Query format:
	// [tag or endpoint type]*.<service>.service.<domain>.
	// <instance_id>.instance.<domain>.
	// For SRV types we also support :
	// _<service>._<tag or endpoint type>.<domain>.

	/// IsDomainName checks if s is a valid domain name
	//  When false is returned the number of labels is not
	// defined.  Also note that this function is extremely liberal; almost any
	// string is a valid domain name as the DNS is 8 bit protocol. It checks if each
	// label fits in 63 characters, but there is no length check for the entire
	// string s. I.e.  a domain name longer than 255 characters is considered valid.
	numberOfLabels, isValidDomain := dns.IsDomainName(question.Name)
	if !isValidDomain {
		response.SetRcode(request, dns.RcodeFormatError)
		return nil, fmt.Errorf("Invalid Domain name %s", question.Name)
	}
	fullDomainRequestArray := dns.SplitDomainName(question.Name)
	if len(fullDomainRequestArray) == 1 || len(fullDomainRequestArray) == 2 {
		response.SetRcode(request, dns.RcodeNameError)
		return nil, fmt.Errorf("service name wasn't included in domain %s", question.Name)
	}
	if fullDomainRequestArray[numberOfLabels-2] == "service" {
		if question.Qtype == dns.TypeSRV && numberOfLabels == 4 &&
			strings.HasPrefix(fullDomainRequestArray[0], "_") &&
			strings.HasPrefix(fullDomainRequestArray[1], "_") {
			// SRV Query :
			tagOrProtocol := fullDomainRequestArray[1][1:]
			serviceName := fullDomainRequestArray[0][1:]
			serviceInstances, err = s.retrieveInstancesForServiceQuery(serviceName, request, response, tagOrProtocol)
		} else {
			serviceName := fullDomainRequestArray[numberOfLabels-3]
			tagsOrProtocol := fullDomainRequestArray[:numberOfLabels-3]
			serviceInstances, err = s.retrieveInstancesForServiceQuery(serviceName, request, response, tagsOrProtocol...)

		}

	} else if fullDomainRequestArray[numberOfLabels-2] == "instance" && (question.Qtype == dns.TypeA ||
		question.Qtype == dns.TypeAAAA) && numberOfLabels == 3 {

		instanceID := fullDomainRequestArray[0]
		serviceInstances, err = s.retrieveInstancesForInstanceQuery(instanceID, request, response)
	}
	return serviceInstances, err
}

func (s *Server) retrieveInstancesForServiceQuery(serviceName string, request, response *dns.Msg, tagOrProtocol ...string) ([]*client.ServiceInstance, error) {
	protocol := ""
	tags := make([]string, 0, len(tagOrProtocol))

	// Split tags and protocol filters
	for _, tag := range tagOrProtocol {
		switch tag {
		case "tcp", "udp", "http", "https":
			if protocol != "" {
				response.SetRcode(request, dns.RcodeFormatError)
				return nil, fmt.Errorf("invalid DNS query: more than one protocol specified")
			}
			protocol = tag
		default:
			tags = append(tags, tag)
		}
	}
	filters := client.InstanceFilter{ServiceName: serviceName, Tags: tags}

	// Dispatch query to registry
	serviceInstances, err := s.config.DiscoveryClient.ListInstances(filters)
	if err != nil {
		response.SetRcode(request, dns.RcodeServerFailure)
		return nil, err
	}

	// Apply protocol filter
	if protocol != "" {
		k := 0
		for _, serviceInstance := range serviceInstances {
			if serviceInstance.Endpoint.Type == protocol {
				serviceInstances[k] = serviceInstance
				k++
			}
		}
		serviceInstances = serviceInstances[:k]
	}

	return serviceInstances, nil
}

func (s *Server) retrieveInstancesForInstanceQuery(instanceID string, request, response *dns.Msg) ([]*client.ServiceInstance, error) {
	serviceInstances, err := s.config.DiscoveryClient.ListInstances(client.InstanceFilter{})
	if err != nil {
		response.SetRcode(request, dns.RcodeServerFailure)
		return serviceInstances, err
	}
	for _, serviceInstance := range serviceInstances {
		if serviceInstance.ID == instanceID {
			return []*client.ServiceInstance{serviceInstance}, nil
		}
	}
	response.SetRcode(request, dns.RcodeNameError)
	return nil, fmt.Errorf("Error : didn't find a service with the id given %s", instanceID)
}

func (s *Server) createRecordsForInstances(question dns.Question, request, response *dns.Msg,
	serviceInstances []*client.ServiceInstance) error {
	numOfMatchingRecords := 0
	for _, serviceInstance := range serviceInstances {
		endPointType := serviceInstance.Endpoint.Type
		var ip net.IP
		var err error
		var port string

		switch endPointType {
		case "tcp", "udp":
			ip, port, err = splitHostPortTCPUDP(serviceInstance.Endpoint.Value)
		case "http", "https":
			ip, port, err = splitHostPortHTTP(serviceInstance.Endpoint.Value)

		default:
			continue
		}
		if err != nil {
			continue
		}
		numOfMatchingRecords++
		if question.Qtype == dns.TypeSRV {

			domainName := s.config.Domain
			instanceID := serviceInstance.ID
			targetName := fmt.Sprintf("%s.instance.%s", instanceID, domainName)
			portNumber, _ := strconv.Atoi(port)
			recordSRV := &dns.SRV{Hdr: dns.RR_Header{
				Name:   question.Name,
				Rrtype: dns.TypeSRV,
				Class:  dns.ClassINET,
				Ttl:    0,
			},
				Port:     uint16(portNumber),
				Priority: 1,
				Target:   targetName,
				Weight:   1,
			}
			response.Answer = append(response.Answer, recordSRV)
			if ip.To4() != nil {
				recordA := createARecord(targetName, ip)
				response.Extra = append(response.Extra, recordA)
			} else if ip.To16() != nil {
				recordAAAA := createAAARecord(targetName, ip)
				response.Extra = append(response.Extra, recordAAAA)
			}

		} else if ip.To4() != nil && question.Qtype == dns.TypeA {
			record := createARecord(question.Name, ip)
			response.Answer = append(response.Answer, record)
		} else if ip.To16() != nil && question.Qtype == dns.TypeAAAA {
			record := createAAARecord(question.Name, ip)
			response.Answer = append(response.Answer, record)
		}
	}
	if numOfMatchingRecords == 0 {
		//Non-Existent Domain
		response.SetRcode(request, dns.RcodeNameError)
		return fmt.Errorf("Non-Existent Domain	 %s", question.Name)

	}
	response.SetRcode(request, dns.RcodeSuccess)
	return nil

}

func splitHostPortTCPUDP(value string) (net.IP, string, error) {
	// Assume value is "host:port"
	host, port, err := net.SplitHostPort(value)

	// Assume value is "host" (no port)
	if err != nil {
		host = value
		port = "0"
	}

	ip := net.ParseIP(host)
	if ip == nil {
		return nil, "", fmt.Errorf("could not parse '%s' as ip:port", value)
	}

	return ip, port, nil
}

func splitHostPortHTTP(value string) (net.IP, string, error) {
	isHTTP := strings.HasPrefix(value, "http://")
	isHTTPS := strings.HasPrefix(value, "https://")
	if !isHTTPS && !isHTTP {
		value = "http://" + value
		isHTTP = true
	}

	parsedURL, err := url.Parse(value)
	if err != nil {
		return nil, "", err
	}

	ip, port, err := splitHostPortTCPUDP(parsedURL.Host)
	if err != nil {
		return nil, "", err
	}

	// Use default port, if not specified
	if port == "0" {
		if isHTTP {
			port = "80"
		} else if isHTTPS {
			port = "443"
		}
	}

	return ip, port, nil

}

func createARecord(questionName string, ip net.IP) *dns.A {
	record := &dns.A{
		Hdr: dns.RR_Header{
			Name:   questionName,
			Rrtype: dns.TypeA,
			Class:  dns.ClassINET,
			Ttl:    0,
		},

		A: ip,
	}
	return record
}

func createAAARecord(questionName string, ip net.IP) *dns.AAAA {
	record := &dns.AAAA{
		Hdr: dns.RR_Header{
			Name:   questionName,
			Rrtype: dns.TypeAAAA,
			Class:  dns.ClassINET,
			Ttl:    0,
		},

		AAAA: ip,
	}
	return record
}

func validate(config *Config) error {
	if config.DiscoveryClient == nil {
		return fmt.Errorf("Discovery client is nil")
	}

	config.Domain = dns.Fqdn(config.Domain)

	return nil
}
<commit_msg>use weight=priority=0 for SRV records<commit_after>// Copyright 2016 IBM Corporation
//
//   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 dns

import (
	"fmt"
	"net"
	"net/url"
	"strconv"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/amalgam8/amalgam8/registry/client"
	"github.com/miekg/dns"
)

// Server represent a DNS server. has config field for port,domain,and client discovery, and the DNS server itself
type Server struct {
	config    Config
	dnsServer *dns.Server
}

// Config represents the DNS server configurations.
type Config struct {
	DiscoveryClient client.Discovery
	Port            uint16
	Domain          string
}

// NewServer creates a new instance of a DNS server with the given configurations
func NewServer(config Config) (*Server, error) {
	err := validate(&config)
	if err != nil {
		return nil, err
	}
	s := &Server{
		config: config,
	}

	// Setup DNS muxing
	mux := dns.NewServeMux()
	mux.HandleFunc(config.Domain, s.handleRequest)

	// Setup a DNS server
	s.dnsServer = &dns.Server{
		Addr:    fmt.Sprintf(":%d", config.Port),
		Net:     "udp",
		Handler: mux,
	}

	return s, nil
}

// ListenAndServe starts the DNS server
func (s *Server) ListenAndServe() error {
	logrus.Info("Starting DNS server")
	err := s.dnsServer.ListenAndServe()

	if err != nil {
		logrus.WithError(err).Errorf("Error starting DNS server")
	}

	return nil
}

// Shutdown stops the DNS server
func (s *Server) Shutdown() error {
	logrus.Info("Shutting down DNS server")
	err := s.dnsServer.Shutdown()

	if err != nil {
		logrus.WithError(err).Errorf("Error shutting down DNS server")
	} else {
		logrus.Info("DNS server has shutdown")
	}

	return err
}

func (s *Server) handleRequest(w dns.ResponseWriter, request *dns.Msg) {
	response := new(dns.Msg)
	response.SetReply(request)
	response.Extra = request.Extra
	response.Authoritative = true
	response.RecursionAvailable = false

	for i, question := range request.Question {
		err := s.handleQuestion(question, request, response)
		if err != nil {
			logrus.WithError(err).Errorf("Error handling DNS question %d: %s", i, question.String())
			// TODO: what should the dns response return ?
			break
		}
	}
	err := w.WriteMsg(response)
	if err != nil {
		logrus.WithError(err).Errorf("Error writing DNS response")
	}
}

func (s *Server) handleQuestion(question dns.Question, request, response *dns.Msg) error {

	switch question.Qclass {
	case dns.ClassINET:
	default:
		response.SetRcode(request, dns.RcodeServerFailure)
		return fmt.Errorf("unsupported DNS question class: %v", dns.Class(question.Qclass).String())
	}

	switch question.Qtype {
	case dns.TypeA:
	case dns.TypeAAAA:
	case dns.TypeSRV:
	default:
		response.SetRcode(request, dns.RcodeServerFailure)
		return fmt.Errorf("unsupported DNS question type: %v", dns.Type(question.Qtype).String())
	}

	serviceInstances, err := s.retrieveServices(question, request, response)

	if err != nil {
		return err
	}
	err = s.createRecordsForInstances(question, request, response, serviceInstances)
	return err

}

func (s *Server) retrieveServices(question dns.Question, request, response *dns.Msg) ([]*client.ServiceInstance, error) {
	var serviceInstances []*client.ServiceInstance
	var err error
	// parse query :
	// Query format:
	// [tag or endpoint type]*.<service>.service.<domain>.
	// <instance_id>.instance.<domain>.
	// For SRV types we also support :
	// _<service>._<tag or endpoint type>.<domain>.

	/// IsDomainName checks if s is a valid domain name
	//  When false is returned the number of labels is not
	// defined.  Also note that this function is extremely liberal; almost any
	// string is a valid domain name as the DNS is 8 bit protocol. It checks if each
	// label fits in 63 characters, but there is no length check for the entire
	// string s. I.e.  a domain name longer than 255 characters is considered valid.
	numberOfLabels, isValidDomain := dns.IsDomainName(question.Name)
	if !isValidDomain {
		response.SetRcode(request, dns.RcodeFormatError)
		return nil, fmt.Errorf("Invalid Domain name %s", question.Name)
	}
	fullDomainRequestArray := dns.SplitDomainName(question.Name)
	if len(fullDomainRequestArray) == 1 || len(fullDomainRequestArray) == 2 {
		response.SetRcode(request, dns.RcodeNameError)
		return nil, fmt.Errorf("service name wasn't included in domain %s", question.Name)
	}
	if fullDomainRequestArray[numberOfLabels-2] == "service" {
		if question.Qtype == dns.TypeSRV && numberOfLabels == 4 &&
			strings.HasPrefix(fullDomainRequestArray[0], "_") &&
			strings.HasPrefix(fullDomainRequestArray[1], "_") {
			// SRV Query :
			tagOrProtocol := fullDomainRequestArray[1][1:]
			serviceName := fullDomainRequestArray[0][1:]
			serviceInstances, err = s.retrieveInstancesForServiceQuery(serviceName, request, response, tagOrProtocol)
		} else {
			serviceName := fullDomainRequestArray[numberOfLabels-3]
			tagsOrProtocol := fullDomainRequestArray[:numberOfLabels-3]
			serviceInstances, err = s.retrieveInstancesForServiceQuery(serviceName, request, response, tagsOrProtocol...)

		}

	} else if fullDomainRequestArray[numberOfLabels-2] == "instance" && (question.Qtype == dns.TypeA ||
		question.Qtype == dns.TypeAAAA) && numberOfLabels == 3 {

		instanceID := fullDomainRequestArray[0]
		serviceInstances, err = s.retrieveInstancesForInstanceQuery(instanceID, request, response)
	}
	return serviceInstances, err
}

func (s *Server) retrieveInstancesForServiceQuery(serviceName string, request, response *dns.Msg, tagOrProtocol ...string) ([]*client.ServiceInstance, error) {
	protocol := ""
	tags := make([]string, 0, len(tagOrProtocol))

	// Split tags and protocol filters
	for _, tag := range tagOrProtocol {
		switch tag {
		case "tcp", "udp", "http", "https":
			if protocol != "" {
				response.SetRcode(request, dns.RcodeFormatError)
				return nil, fmt.Errorf("invalid DNS query: more than one protocol specified")
			}
			protocol = tag
		default:
			tags = append(tags, tag)
		}
	}
	filters := client.InstanceFilter{ServiceName: serviceName, Tags: tags}

	// Dispatch query to registry
	serviceInstances, err := s.config.DiscoveryClient.ListInstances(filters)
	if err != nil {
		response.SetRcode(request, dns.RcodeServerFailure)
		return nil, err
	}

	// Apply protocol filter
	if protocol != "" {
		k := 0
		for _, serviceInstance := range serviceInstances {
			if serviceInstance.Endpoint.Type == protocol {
				serviceInstances[k] = serviceInstance
				k++
			}
		}
		serviceInstances = serviceInstances[:k]
	}

	return serviceInstances, nil
}

func (s *Server) retrieveInstancesForInstanceQuery(instanceID string, request, response *dns.Msg) ([]*client.ServiceInstance, error) {
	serviceInstances, err := s.config.DiscoveryClient.ListInstances(client.InstanceFilter{})
	if err != nil {
		response.SetRcode(request, dns.RcodeServerFailure)
		return serviceInstances, err
	}
	for _, serviceInstance := range serviceInstances {
		if serviceInstance.ID == instanceID {
			return []*client.ServiceInstance{serviceInstance}, nil
		}
	}
	response.SetRcode(request, dns.RcodeNameError)
	return nil, fmt.Errorf("Error : didn't find a service with the id given %s", instanceID)
}

func (s *Server) createRecordsForInstances(question dns.Question, request, response *dns.Msg,
	serviceInstances []*client.ServiceInstance) error {
	numOfMatchingRecords := 0
	for _, serviceInstance := range serviceInstances {
		endPointType := serviceInstance.Endpoint.Type
		var ip net.IP
		var err error
		var port string

		switch endPointType {
		case "tcp", "udp":
			ip, port, err = splitHostPortTCPUDP(serviceInstance.Endpoint.Value)
		case "http", "https":
			ip, port, err = splitHostPortHTTP(serviceInstance.Endpoint.Value)

		default:
			continue
		}
		if err != nil {
			continue
		}
		numOfMatchingRecords++
		if question.Qtype == dns.TypeSRV {

			domainName := s.config.Domain
			instanceID := serviceInstance.ID
			targetName := fmt.Sprintf("%s.instance.%s", instanceID, domainName)
			portNumber, _ := strconv.Atoi(port)
			recordSRV := createSRVRecord(question.Name, portNumber, targetName)
			response.Answer = append(response.Answer, recordSRV)
			if ip.To4() != nil {
				recordA := createARecord(targetName, ip)
				response.Extra = append(response.Extra, recordA)
			} else if ip.To16() != nil {
				recordAAAA := createAAAARecord(targetName, ip)
				response.Extra = append(response.Extra, recordAAAA)
			}

		} else if ip.To4() != nil && question.Qtype == dns.TypeA {
			record := createARecord(question.Name, ip)
			response.Answer = append(response.Answer, record)
		} else if ip.To16() != nil && question.Qtype == dns.TypeAAAA {
			record := createAAAARecord(question.Name, ip)
			response.Answer = append(response.Answer, record)
		}
	}
	if numOfMatchingRecords == 0 {
		//Non-Existent Domain
		response.SetRcode(request, dns.RcodeNameError)
		return fmt.Errorf("Non-Existent Domain	 %s", question.Name)

	}
	response.SetRcode(request, dns.RcodeSuccess)
	return nil

}

func splitHostPortTCPUDP(value string) (net.IP, string, error) {
	// Assume value is "host:port"
	host, port, err := net.SplitHostPort(value)

	// Assume value is "host" (no port)
	if err != nil {
		host = value
		port = "0"
	}

	ip := net.ParseIP(host)
	if ip == nil {
		return nil, "", fmt.Errorf("could not parse '%s' as ip:port", value)
	}

	return ip, port, nil
}

func splitHostPortHTTP(value string) (net.IP, string, error) {
	isHTTP := strings.HasPrefix(value, "http://")
	isHTTPS := strings.HasPrefix(value, "https://")
	if !isHTTPS && !isHTTP {
		value = "http://" + value
		isHTTP = true
	}

	parsedURL, err := url.Parse(value)
	if err != nil {
		return nil, "", err
	}

	ip, port, err := splitHostPortTCPUDP(parsedURL.Host)
	if err != nil {
		return nil, "", err
	}

	// Use default port, if not specified
	if port == "0" {
		if isHTTP {
			port = "80"
		} else if isHTTPS {
			port = "443"
		}
	}

	return ip, port, nil

}

func createARecord(name string, ip net.IP) *dns.A {
	record := &dns.A{
		Hdr: dns.RR_Header{
			Name:   name,
			Rrtype: dns.TypeA,
			Class:  dns.ClassINET,
			Ttl:    0,
		},
		A: ip,
	}
	return record
}

func createAAAARecord(name string, ip net.IP) *dns.AAAA {
	record := &dns.AAAA{
		Hdr: dns.RR_Header{
			Name:   name,
			Rrtype: dns.TypeAAAA,
			Class:  dns.ClassINET,
			Ttl:    0,
		},
		AAAA: ip,
	}
	return record
}

func createSRVRecord(name string, port int, target string) *dns.AAAA {
	record := &dns.SRV{
		Hdr: dns.RR_Header{
			Name:   name,
			Rrtype: dns.TypeSRV,
			Class:  dns.ClassINET,
			Ttl:    0,
		},
		Port:     uint16(port),
		Priority: 0,
		Weight:   0,
		Target:   target,
	}
	return record
}

func validate(config *Config) error {
	if config.DiscoveryClient == nil {
		return fmt.Errorf("Discovery client is nil")
	}

	config.Domain = dns.Fqdn(config.Domain)

	return nil
}
<|endoftext|>