<commit_before>package vsphere

import (
	"fmt"
	"log"
	"net"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/find"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/property"
	"github.com/vmware/govmomi/vim25/mo"
	"golang.org/x/net/context"
)

var DefaultDNSSuffixes = []string{
	"vsphere.local",
}

var DefaultDNSServers = []string{
	"8.8.8.8",
	"8.8.4.4",
}

func resourceVSphereVirtualMachine() *schema.Resource {
	return &schema.Resource{
		Create: resourceVSphereVirtualMachineCreate,
		Read:   resourceVSphereVirtualMachineRead,
		Update: resourceVSphereVirtualMachineUpdate,
		Delete: resourceVSphereVirtualMachineDelete,

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

			"vcpu": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: false,
			},

			"memory": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: false,
			},

			"datacenter": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
			},

			"cluster": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
			},

			"resource_pool": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
			},

			"gateway": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
			},

			"domain": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
				Default:  "vsphere.local",
			},

			"time_zone": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
				Default:  "Etc/UTC",
			},

			"dns_suffix": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				ForceNew: false,
			},

			"dns_server": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				ForceNew: false,
			},

			"network_interface": &schema.Schema{
				Type:     schema.TypeList,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"label": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
							ForceNew: false,
						},

						"ip_address": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
							ForceNew: false,
						},

						"subnet_mask": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
							ForceNew: false,
						},

						"adapter_type": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: false,
						},
					},
				},
			},

			"disk": &schema.Schema{
				Type:     schema.TypeList,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"template": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: false,
						},

						"datastore": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: false,
						},

						"size": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
							ForceNew: false,
						},

						"iops": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
							ForceNew: false,
						},
					},
				},
			},
		},
	}
}

func resourceVSphereVirtualMachineCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*govmomi.Client)

	vm := virtualMachine{
		name:     d.Get("name").(string),
		vcpu:     d.Get("vcpu").(int),
		memoryMb: int64(d.Get("memory").(int)),
	}

	if v, ok := d.GetOk("datacenter"); ok {
		vm.datacenter = v.(string)
	}

	if v, ok := d.GetOk("cluster"); ok {
		vm.cluster = v.(string)
	}

	if v, ok := d.GetOk("resource_pool"); ok {
		vm.resourcePool = v.(string)
	}

	if v, ok := d.GetOk("gateway"); ok {
		vm.gateway = v.(string)
	}

	if v, ok := d.GetOk("domain"); ok {
		vm.domain = v.(string)
	}

	if v, ok := d.GetOk("time_zone"); ok {
		vm.timeZone = v.(string)
	}

	dns_suffix := d.Get("dns_suffix.#").(int)
	if dns_suffix > 0 {
		vm.dnsSuffixes = make([]string, 0, dns_suffix)
		for i := 0; i < dns_suffix; i++ {
			s := fmt.Sprintf("dns_suffix.%d", i)
			vm.dnsSuffixes = append(vm.dnsSuffixes, d.Get(s).(string))
		}
	} else {
		vm.dnsSuffixes = DefaultDNSSuffixes
	}

	dns_server := d.Get("dns_server.#").(int)
	if dns_server > 0 {
		vm.dnsServers = make([]string, 0, dns_server)
		for i := 0; i < dns_server; i++ {
			s := fmt.Sprintf("dns_server.%d", i)
			vm.dnsServers = append(vm.dnsServers, d.Get(s).(string))
		}
	} else {
		vm.dnsServers = DefaultDNSServers
	}

	networksCount := d.Get("network_interface.#").(int)
	networks := make([]networkInterface, networksCount)
	for i := 0; i < networksCount; i++ {
		prefix := fmt.Sprintf("network_interface.%d", i)
		networks[i].label = d.Get(prefix + ".label").(string)
		if v := d.Get(prefix + ".ip_address"); v != nil {
			networks[i].ipAddress = d.Get(prefix + ".ip_address").(string)
		}
		if v := d.Get(prefix + ".subnet_mask"); v != nil {
			networks[i].subnetMask = d.Get(prefix + ".subnet_mask").(string)
		}
	}
	vm.networkInterfaces = networks
	log.Printf("[DEBUG] network_interface init: %v", networks)

	diskCount := d.Get("disk.#").(int)
	disks := make([]hardDisk, diskCount)
	for i := 0; i < diskCount; i++ {
		prefix := fmt.Sprintf("disk.%d", i)
		if i == 0 {
			if v := d.Get(prefix + ".template"); v != "" {
				vm.template = d.Get(prefix + ".template").(string)
			} else {
				if v := d.Get(prefix + ".size"); v != "" {
					disks[i].size = int64(d.Get(prefix + ".size").(int))
				} else {
					return fmt.Errorf("If template argument is not specified, size argument is required.")
				}
			}
			if v := d.Get(prefix + ".datastore"); v != "" {
				vm.datastore = d.Get(prefix + ".datastore").(string)
			}
		} else {
			if v := d.Get(prefix + ".size"); v != "" {
				disks[i].size = int64(d.Get(prefix + ".size").(int))
			} else {
				return fmt.Errorf("Size argument is required.")
			}
		}
		if v := d.Get(prefix + ".iops"); v != "" {
			disks[i].iops = int64(d.Get(prefix + ".iops").(int))
		}
	}
	vm.hardDisks = disks
	log.Printf("[DEBUG] disk init: %v", disks)

	if vm.template != "" {
		err := vm.deployVirtualMachine(client)
		if err != nil {
			return fmt.Errorf("error: %s", err)
		}
	} else {
		err := vm.createVirtualMachine(client)
		if err != nil {
			return fmt.Errorf("error: %s", err)
		}
	}
	d.SetId(vm.name)
	log.Printf("[INFO] Created virtual machine: %s", d.Id())

	return resourceVSphereVirtualMachineRead(d, meta)
}

func resourceVSphereVirtualMachineRead(d *schema.ResourceData, meta interface{}) error {
	var dc *object.Datacenter
	var err error

	client := meta.(*govmomi.Client)
	finder := find.NewFinder(client.Client, true)

	if v, ok := d.GetOk("datacenter"); ok {
		dc, err = finder.Datacenter(context.TODO(), v.(string))
		if err != nil {
			return err
		}
	} else {
		dc, err = finder.DefaultDatacenter(context.TODO())
		if err != nil {
			return err
		}
	}

	finder = finder.SetDatacenter(dc)
	vm, err := finder.VirtualMachine(context.TODO(), d.Get("name").(string))
	if err != nil {
		log.Printf("[ERROR] Virtual machine not found: %s", d.Get("name").(string))
		d.SetId("")
		return nil
	}

	var mvm mo.VirtualMachine

	collector := property.DefaultCollector(client.Client)
	if err := collector.RetrieveOne(context.TODO(), vm.Reference(), []string{"guest", "summary", "datastore"}, &mvm); err != nil {
		log.Printf("[ERROR] %#v", err)
	}

	log.Printf("[DEBUG] %#v", dc)
	log.Printf("[DEBUG] %#v", mvm.Summary.Config)
	log.Printf("[DEBUG] %#v", mvm.Guest.Net)

	networkInterfaces := make([]map[string]interface{}, len(mvm.Guest.Net))
	for i, v := range mvm.Guest.Net {
		log.Printf("[DEBUG] %#v", v.Network)
		networkInterfaces[i] = make(map[string]interface{})
		networkInterfaces[i]["label"] = v.Network
		if len(v.IpAddress) > 0 {
			log.Printf("[DEBUG] %#v", v.IpAddress[0])
			networkInterfaces[i]["ip_address"] = v.IpAddress[0]

			m := net.CIDRMask(v.IpConfig.IpAddress[0].PrefixLength, 32)
			subnetMask := net.IPv4(m[0], m[1], m[2], m[3])
			networkInterfaces[i]["subnet_mask"] = subnetMask.String()
			log.Printf("[DEBUG] %#v", subnetMask.String())
		}
	}
	d.Set("network_interface", networkInterfaces)

	var rootDatastore string
	for _, v := range mvm.Datastore {
		var md mo.Datastore
		if err := collector.RetrieveOne(context.TODO(), v, []string{"name", "parent"}, &md); err != nil {
			log.Printf("[ERROR] %#v", err)
		}
		if md.Parent.Type == "StoragePod" {
			var msp mo.StoragePod
			if err := collector.RetrieveOne(context.TODO(), *md.Parent, []string{"name"}, &msp); err != nil {
				log.Printf("[ERROR] %#v", err)
			}
			rootDatastore = msp.Name
			log.Printf("[DEBUG] %#v", msp.Name)
		} else {
			rootDatastore = md.Name
			log.Printf("[DEBUG] %#v", md.Name)
		}
		break
	}

	d.Set("datacenter", dc)
	d.Set("memory", mvm.Summary.Config.MemorySizeMB)
	d.Set("cpu", mvm.Summary.Config.NumCpu)
	d.Set("datastore", rootDatastore)

	// Initialize the connection info
	d.SetConnInfo(map[string]string{
		"type": "ssh",
		"host": networkInterfaces[0]["ip_address"].(string),
	})

	return nil
}

func resourceVSphereVirtualMachineUpdate(d *schema.ResourceData, meta interface{}) error {
	return nil
}

func resourceVSphereVirtualMachineDelete(d *schema.ResourceData, meta interface{}) error {
	var dc *object.Datacenter
	var err error

	client := meta.(*govmomi.Client)
	finder := find.NewFinder(client.Client, true)

	if v, ok := d.GetOk("datacenter"); ok {
		dc, err = finder.Datacenter(context.TODO(), v.(string))
		if err != nil {
			return err
		}
	} else {
		dc, err = finder.DefaultDatacenter(context.TODO())
		if err != nil {
			return err
		}
	}

	finder = finder.SetDatacenter(dc)
	vm, err := finder.VirtualMachine(context.TODO(), d.Get("name").(string))
	if err != nil {
		return err
	}

	log.Printf("[INFO] Deleting virtual machine: %s", d.Id())

	task, err := vm.PowerOff(context.TODO())
	if err != nil {
		return err
	}

	err = task.Wait(context.TODO())
	if err != nil {
		return err
	}

	task, err = vm.Destroy(context.TODO())
	if err != nil {
		return err
	}

	err = task.Wait(context.TODO())
	if err != nil {
		return err
	}

	d.SetId("")
	return nil
}
<commit_msg>Add boot delay feature<commit_after>package vsphere

import (
	"fmt"
	"log"
	"net"
	"time"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/find"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/property"
	"github.com/vmware/govmomi/vim25/mo"
	"golang.org/x/net/context"
)

var DefaultDNSSuffixes = []string{
	"vsphere.local",
}

var DefaultDNSServers = []string{
	"8.8.8.8",
	"8.8.4.4",
}

func resourceVSphereVirtualMachine() *schema.Resource {
	return &schema.Resource{
		Create: resourceVSphereVirtualMachineCreate,
		Read:   resourceVSphereVirtualMachineRead,
		Update: resourceVSphereVirtualMachineUpdate,
		Delete: resourceVSphereVirtualMachineDelete,

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

			"vcpu": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: false,
			},

			"memory": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: false,
			},

			"datacenter": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
			},

			"cluster": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
			},

			"resource_pool": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
			},

			"gateway": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
			},

			"domain": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
				Default:  "vsphere.local",
			},

			"time_zone": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: false,
				Default:  "Etc/UTC",
			},

			"dns_suffix": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				ForceNew: false,
			},

			"dns_server": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				ForceNew: false,
			},

			"network_interface": &schema.Schema{
				Type:     schema.TypeList,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"label": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
							ForceNew: false,
						},

						"ip_address": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
							ForceNew: false,
						},

						"subnet_mask": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
							ForceNew: false,
						},

						"adapter_type": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: false,
						},
					},
				},
			},

			"disk": &schema.Schema{
				Type:     schema.TypeList,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"template": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: false,
						},

						"datastore": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: false,
						},

						"size": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
							ForceNew: false,
						},

						"iops": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
							ForceNew: false,
						},
					},
				},
			},

			"boot_delay": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
		},
	}
}

func resourceVSphereVirtualMachineCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*govmomi.Client)

	vm := virtualMachine{
		name:     d.Get("name").(string),
		vcpu:     d.Get("vcpu").(int),
		memoryMb: int64(d.Get("memory").(int)),
	}

	if v, ok := d.GetOk("datacenter"); ok {
		vm.datacenter = v.(string)
	}

	if v, ok := d.GetOk("cluster"); ok {
		vm.cluster = v.(string)
	}

	if v, ok := d.GetOk("resource_pool"); ok {
		vm.resourcePool = v.(string)
	}

	if v, ok := d.GetOk("gateway"); ok {
		vm.gateway = v.(string)
	}

	if v, ok := d.GetOk("domain"); ok {
		vm.domain = v.(string)
	}

	if v, ok := d.GetOk("time_zone"); ok {
		vm.timeZone = v.(string)
	}

	dns_suffix := d.Get("dns_suffix.#").(int)
	if dns_suffix > 0 {
		vm.dnsSuffixes = make([]string, 0, dns_suffix)
		for i := 0; i < dns_suffix; i++ {
			s := fmt.Sprintf("dns_suffix.%d", i)
			vm.dnsSuffixes = append(vm.dnsSuffixes, d.Get(s).(string))
		}
	} else {
		vm.dnsSuffixes = DefaultDNSSuffixes
	}

	dns_server := d.Get("dns_server.#").(int)
	if dns_server > 0 {
		vm.dnsServers = make([]string, 0, dns_server)
		for i := 0; i < dns_server; i++ {
			s := fmt.Sprintf("dns_server.%d", i)
			vm.dnsServers = append(vm.dnsServers, d.Get(s).(string))
		}
	} else {
		vm.dnsServers = DefaultDNSServers
	}

	networksCount := d.Get("network_interface.#").(int)
	networks := make([]networkInterface, networksCount)
	for i := 0; i < networksCount; i++ {
		prefix := fmt.Sprintf("network_interface.%d", i)
		networks[i].label = d.Get(prefix + ".label").(string)
		if v, ok := d.GetOk(prefix + ".ip_address"); ok {
			networks[i].ipAddress = v.(string)
		}
		if v, ok := d.GetOk(prefix + ".subnet_mask"); ok {
			networks[i].subnetMask = v.(string)
		}
	}
	vm.networkInterfaces = networks
	log.Printf("[DEBUG] network_interface init: %v", networks)

	diskCount := d.Get("disk.#").(int)
	disks := make([]hardDisk, diskCount)
	for i := 0; i < diskCount; i++ {
		prefix := fmt.Sprintf("disk.%d", i)
		if i == 0 {
			if v, ok := d.GetOk(prefix + ".template"); ok {
				vm.template = v.(string)
			} else {
				if v, ok := d.GetOk(prefix + ".size"); ok {
					disks[i].size = int64(v.(int))
				} else {
					return fmt.Errorf("If template argument is not specified, size argument is required.")
				}
			}
			if v, ok := d.GetOk(prefix + ".datastore"); ok {
				vm.datastore = v.(string)
			}
		} else {
			if v, ok := d.GetOk(prefix + ".size"); ok {
				disks[i].size = int64(v.(int))
			} else {
				return fmt.Errorf("Size argument is required.")
			}
		}
		if v, ok := d.GetOk(prefix + ".iops"); ok {
			disks[i].iops = int64(v.(int))
		}
	}
	vm.hardDisks = disks
	log.Printf("[DEBUG] disk init: %v", disks)

	if vm.template != "" {
		err := vm.deployVirtualMachine(client)
		if err != nil {
			return fmt.Errorf("error: %s", err)
		}
	} else {
		err := vm.createVirtualMachine(client)
		if err != nil {
			return fmt.Errorf("error: %s", err)
		}
	}

	if _, ok := d.GetOk("network_interface.0.ip_address"); !ok {
		if v, ok := d.GetOk("boot_delay"); ok {
			var dc *object.Datacenter
			var err error
			finder := find.NewFinder(client.Client, true)

			if v, ok := d.GetOk("datacenter"); ok {
				dc, err = finder.Datacenter(context.TODO(), v.(string))
				if err != nil {
					return err
				}
			} else {
				dc, err = finder.DefaultDatacenter(context.TODO())
				if err != nil {
					return err
				}
			}

			finder = finder.SetDatacenter(dc)
			vm, err := finder.VirtualMachine(context.TODO(), d.Get("name").(string))
			if err != nil {
				log.Printf("[ERROR] Virtual machine not found: %s", d.Get("name").(string))
				d.SetId("")
				return nil
			}

			var mvm mo.VirtualMachine
			collector := property.DefaultCollector(client.Client)
			if err := collector.RetrieveOne(context.TODO(), vm.Reference(), []string{"summary"}, &mvm); err != nil {
				log.Printf("[ERROR] %#v", err)
			}
			bootTime := *mvm.Summary.Runtime.BootTime
			duration := time.Since(bootTime)
			remainingBootDelay := float64(v.(int)) - float64(duration.Seconds())
			if remainingBootDelay > 0 {
				time.Sleep(time.Duration(int(remainingBootDelay)))
			}

			ip, err := vm.WaitForIP(context.TODO())
			if err != nil {
				log.Printf("[ERROR] Wait for IP function error, %s", err)
				d.SetId("")
				return nil
			}
			log.Printf("[DEBUG] new ip address after boot delay: %v", ip)
		}
	}
	d.SetId(vm.name)
	log.Printf("[INFO] Created virtual machine: %s", d.Id())

	return resourceVSphereVirtualMachineRead(d, meta)
}

func resourceVSphereVirtualMachineRead(d *schema.ResourceData, meta interface{}) error {
	var dc *object.Datacenter
	var err error

	client := meta.(*govmomi.Client)
	finder := find.NewFinder(client.Client, true)

	if v, ok := d.GetOk("datacenter"); ok {
		dc, err = finder.Datacenter(context.TODO(), v.(string))
		if err != nil {
			return err
		}
	} else {
		dc, err = finder.DefaultDatacenter(context.TODO())
		if err != nil {
			return err
		}
	}

	finder = finder.SetDatacenter(dc)
	vm, err := finder.VirtualMachine(context.TODO(), d.Get("name").(string))
	if err != nil {
		log.Printf("[ERROR] Virtual machine not found: %s", d.Get("name").(string))
		d.SetId("")
		return nil
	}

	var mvm mo.VirtualMachine

	collector := property.DefaultCollector(client.Client)
	if err := collector.RetrieveOne(context.TODO(), vm.Reference(), []string{"guest", "summary", "datastore"}, &mvm); err != nil {
		log.Printf("[ERROR] %#v", err)
	}

	log.Printf("[DEBUG] %#v", dc)
	log.Printf("[DEBUG] %#v", mvm.Summary.Config)
	log.Printf("[DEBUG] %#v", mvm.Guest.Net)

	networkInterfaces := make([]map[string]interface{}, len(mvm.Guest.Net))
	for i, v := range mvm.Guest.Net {
		log.Printf("[DEBUG] %#v", v.Network)
		networkInterfaces[i] = make(map[string]interface{})
		networkInterfaces[i]["label"] = v.Network
		if len(v.IpAddress) > 0 {
			log.Printf("[DEBUG] %#v", v.IpAddress[0])
			networkInterfaces[i]["ip_address"] = v.IpAddress[0]

			m := net.CIDRMask(v.IpConfig.IpAddress[0].PrefixLength, 32)
			subnetMask := net.IPv4(m[0], m[1], m[2], m[3])
			networkInterfaces[i]["subnet_mask"] = subnetMask.String()
			log.Printf("[DEBUG] %#v", subnetMask.String())
		}
	}
	d.Set("network_interface", networkInterfaces)

	var rootDatastore string
	for _, v := range mvm.Datastore {
		var md mo.Datastore
		if err := collector.RetrieveOne(context.TODO(), v, []string{"name", "parent"}, &md); err != nil {
			log.Printf("[ERROR] %#v", err)
		}
		if md.Parent.Type == "StoragePod" {
			var msp mo.StoragePod
			if err := collector.RetrieveOne(context.TODO(), *md.Parent, []string{"name"}, &msp); err != nil {
				log.Printf("[ERROR] %#v", err)
			}
			rootDatastore = msp.Name
			log.Printf("[DEBUG] %#v", msp.Name)
		} else {
			rootDatastore = md.Name
			log.Printf("[DEBUG] %#v", md.Name)
		}
		break
	}

	d.Set("datacenter", dc)
	d.Set("memory", mvm.Summary.Config.MemorySizeMB)
	d.Set("cpu", mvm.Summary.Config.NumCpu)
	d.Set("datastore", rootDatastore)

	// Initialize the connection info
	d.SetConnInfo(map[string]string{
		"type": "ssh",
		"host": networkInterfaces[0]["ip_address"].(string),
	})

	return nil
}

func resourceVSphereVirtualMachineUpdate(d *schema.ResourceData, meta interface{}) error {
	return nil
}

func resourceVSphereVirtualMachineDelete(d *schema.ResourceData, meta interface{}) error {
	var dc *object.Datacenter
	var err error

	client := meta.(*govmomi.Client)
	finder := find.NewFinder(client.Client, true)

	if v, ok := d.GetOk("datacenter"); ok {
		dc, err = finder.Datacenter(context.TODO(), v.(string))
		if err != nil {
			return err
		}
	} else {
		dc, err = finder.DefaultDatacenter(context.TODO())
		if err != nil {
			return err
		}
	}

	finder = finder.SetDatacenter(dc)
	vm, err := finder.VirtualMachine(context.TODO(), d.Get("name").(string))
	if err != nil {
		return err
	}

	log.Printf("[INFO] Deleting virtual machine: %s", d.Id())

	task, err := vm.PowerOff(context.TODO())
	if err != nil {
		return err
	}

	err = task.Wait(context.TODO())
	if err != nil {
		return err
	}

	task, err = vm.Destroy(context.TODO())
	if err != nil {
		return err
	}

	err = task.Wait(context.TODO())
	if err != nil {
		return err
	}

	d.SetId("")
	return nil
}
<|endoftext|>


<commit_before>package vision

import (
	"os"
	"testing"

	"github.com/daviddengcn/go-assert"
)

func TestLoad(t *testing.T) {
	datasets := [...]struct {
		width, height int
		fn            string
	}{
		{177, 177, "gray-177x177.jpg"},
		{500, 374, "ycbcr-500x374.jpg"},
		{500, 374, "pal-500x374.gif"},
		{16, 16, "nrgba-16x16.png"},
	}

	os.MkdirAll("testout", 0755)

	for _, d := range datasets {
		img, err := ImageFromFile("testdata/" + d.fn)
		if err != nil {
			t.Error(err)
			continue
		}
		t.Logf("Image loaded: %v", d.fn)

		var m GrayImage
		m.SetImage(img)

		if m.Width != d.width || m.Height != d.height {
			t.Errorf("(GrayImage) Wrong size: %d, %d, expected %d, %d",
				m.Width, m.Height, d.width, d.height)
		}
		outImg := m.AsImage()
		if outImg.Bounds().Dx() != d.width || outImg.Bounds().Dy() != d.height {
			t.Errorf("(AsImage) Wrong size: %d, %d, expected %d, %d",
				outImg.Bounds().Dx(), outImg.Bounds().Dy(), d.width, d.height)
		}

		if err := SaveImageAsPng(outImg, "testout/gray-"+d.fn+".png"); err != nil {
			t.Error(err)
			continue
		}

		var rgbImg RGBImage
		rgbImg.SetImage(img)
		if rgbImg.Width != d.width || rgbImg.Height != d.height {
			t.Errorf("(RGBImage) Wrong size: %d, %d, expected %d, %d",
				rgbImg.Width, rgbImg.Height, d.width, d.height)
		}

		outImg = rgbImg.AsImage()
		if outImg.Bounds().Dx() != d.width || outImg.Bounds().Dy() != d.height {
			t.Errorf("(AsImage) Wrong size: %d, %d, expected %d, %d",
				outImg.Bounds().Dx(), outImg.Bounds().Dy(), d.width, d.height)
		}

		if err := SaveImageAsPng(outImg, "testout/rgb-"+d.fn+".png"); err != nil {
			t.Error(err)
			continue
		}
	}
}

func TestResize(t *testing.T) {
	var gray GrayImage
	gray.Resize(Size{123, 456})
	assert.Equals(t, "Size", gray.Size, Size{123, 456})
	assert.Equals(t, "len(Pixels)", len(gray.Pixels), 123*456)

	var ig IntGrayImage
	ig.Resize(Size{123, 555})
	assert.Equals(t, "Size", ig.Size, Size{123, 555})
	assert.Equals(t, "len(Pixels)", len(ig.Pixels), 123*555)

	var rgb RGBImage
	rgb.Resize(Size{123, 456})
	assert.Equals(t, "Size", rgb.Size, Size{123, 456})
	assert.Equals(t, "len(Pixels)", len(rgb.Pixels), 123*456)
}

func TestFill(t *testing.T) {
	var gray GrayImage
	gray.Resize(Size{456, 123})
	gray.Fill(123)
	assert.Equals(t, "[0,0]", gray.Pixels[0], byte(123))

	var ig IntGrayImage
	ig.Resize(Size{555, 123})
	ig.Fill(12345)
	assert.Equals(t, "[0,0]", ig.Pixels[0], int(12345))

	var rgb RGBImage
	rgb.Resize(Size{456, 123})
	rgb.Fill(RGB{1, 2, 3})
	assert.Equals(t, "[0,0]", rgb.Pixels[0], RGB{1, 2, 3})
}
<commit_msg>Change test package<commit_after>package vision

import (
	"os"
	"testing"

	"github.com/golangplus/testing/assert"
)

func TestLoad(t *testing.T) {
	datasets := [...]struct {
		width, height int
		fn            string
	}{
		{177, 177, "gray-177x177.jpg"},
		{500, 374, "ycbcr-500x374.jpg"},
		{500, 374, "pal-500x374.gif"},
		{16, 16, "nrgba-16x16.png"},
	}

	os.MkdirAll("testout", 0755)

	for _, d := range datasets {
		img, err := ImageFromFile("testdata/" + d.fn)
		if err != nil {
			t.Error(err)
			continue
		}
		t.Logf("Image loaded: %v", d.fn)

		var m GrayImage
		m.SetImage(img)

		if m.Width != d.width || m.Height != d.height {
			t.Errorf("(GrayImage) Wrong size: %d, %d, expected %d, %d",
				m.Width, m.Height, d.width, d.height)
		}
		outImg := m.AsImage()
		if outImg.Bounds().Dx() != d.width || outImg.Bounds().Dy() != d.height {
			t.Errorf("(AsImage) Wrong size: %d, %d, expected %d, %d",
				outImg.Bounds().Dx(), outImg.Bounds().Dy(), d.width, d.height)
		}

		if err := SaveImageAsPng(outImg, "testout/gray-"+d.fn+".png"); err != nil {
			t.Error(err)
			continue
		}

		var rgbImg RGBImage
		rgbImg.SetImage(img)
		if rgbImg.Width != d.width || rgbImg.Height != d.height {
			t.Errorf("(RGBImage) Wrong size: %d, %d, expected %d, %d",
				rgbImg.Width, rgbImg.Height, d.width, d.height)
		}

		outImg = rgbImg.AsImage()
		if outImg.Bounds().Dx() != d.width || outImg.Bounds().Dy() != d.height {
			t.Errorf("(AsImage) Wrong size: %d, %d, expected %d, %d",
				outImg.Bounds().Dx(), outImg.Bounds().Dy(), d.width, d.height)
		}

		if err := SaveImageAsPng(outImg, "testout/rgb-"+d.fn+".png"); err != nil {
			t.Error(err)
			continue
		}
	}
}

func TestResize(t *testing.T) {
	var gray GrayImage
	gray.Resize(Size{123, 456})
	assert.Equal(t, "Size", gray.Size, Size{123, 456})
	assert.Equal(t, "len(Pixels)", len(gray.Pixels), 123*456)

	var ig IntGrayImage
	ig.Resize(Size{123, 555})
	assert.Equal(t, "Size", ig.Size, Size{123, 555})
	assert.Equal(t, "len(Pixels)", len(ig.Pixels), 123*555)

	var rgb RGBImage
	rgb.Resize(Size{123, 456})
	assert.Equal(t, "Size", rgb.Size, Size{123, 456})
	assert.Equal(t, "len(Pixels)", len(rgb.Pixels), 123*456)
}

func TestFill(t *testing.T) {
	var gray GrayImage
	gray.Resize(Size{456, 123})
	gray.Fill(123)
	assert.Equal(t, "[0,0]", gray.Pixels[0], byte(123))

	var ig IntGrayImage
	ig.Resize(Size{555, 123})
	ig.Fill(12345)
	assert.Equal(t, "[0,0]", ig.Pixels[0], int(12345))

	var rgb RGBImage
	rgb.Resize(Size{456, 123})
	rgb.Fill(RGB{1, 2, 3})
	assert.Equal(t, "[0,0]", rgb.Pixels[0], RGB{1, 2, 3})
}
<|endoftext|>


<commit_before>// Copyright 2018 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package model

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	authn "istio.io/api/authentication/v1alpha1"
	"istio.io/istio/pkg/cache"
	"istio.io/istio/pkg/log"
)

const (
	// https://openid.net/specs/openid-connect-discovery-1_0.html
	// OpenID Providers supporting Discovery MUST make a JSON document available at the path
	// formed by concatenating the string /.well-known/openid-configuration to the Issuer.
	openIDDiscoveryCfgURLSuffix = "/.well-known/openid-configuration"

	// OpenID Discovery web request timeout.
	jwksHTTPTimeOutInSec = 5

	// JwksURI Cache expiration time duration, individual cached JwksURI item will be removed
	// from cache after its duration expires.
	jwksURICacheExpiration = time.Hour * 24

	// JwksURI Cache eviction time duration, cache eviction is done on a periodic basis,
	// jwksURICacheEviction specifies the frequency at which eviction activities take place.
	jwksURICacheEviction = time.Minute * 30

	// JwtPubKeyExpireDuration is the expire duration for JWT public key in the cache.
	// After this duration expire, refresher job will fetch key for the cached item again.
	JwtPubKeyExpireDuration = time.Hour

	// JwtPubKeyEvictionDuration is the life duration for cached item.
	// Cached item will be removed from the cache if it hasn't been used longer than JwtPubKeyEvictionDuration.
	JwtPubKeyEvictionDuration = 24 * 7 * time.Hour

	// JwtPubKeyRefreshInterval is the running interval of JWT pubKey refresh job.
	JwtPubKeyRefreshInterval = time.Minute * 20

	// PublicRootCABundlePath is the path of public root CA bundle in pilot container.
	publicRootCABundlePath = "/cacert.pem"
)

// jwtPubKeyEntry is a single cached entry for jwt public key.
type jwtPubKeyEntry struct {
	pubKey string

	// Cached item will be fetched again by refresher job if (time.now >= expireTime).
	expireTime time.Time

	// Cached item's last used time, which is set in GetPublicKey.
	lastUsedTime time.Time
}

// jwksResolver is resolver for jwksURI and jwt public key.
type jwksResolver struct {
	// cache for jwksURI.
	JwksURICache cache.ExpiringCache

	// cache for JWT public key.
	// map key is jwksURI, map value is jwtPubKeyEntry.
	keyEntries sync.Map

	secureHTTPClient *http.Client
	httpClient       *http.Client
	closing          chan bool
	refreshTicker    *time.Ticker

	expireDuration time.Duration

	// Cached key will be removed from cache if (time.now - cachedItem.lastUsedTime >= evictionDuration), this prevents key cache growing indefinitely.
	evictionDuration time.Duration

	// Refresher job running interval.
	refreshInterval time.Duration

	// How may times refresh job has detected JWT public key change happened, used in unit test.
	keyChangedCount uint64
}

// newJwksResolver creates new instance of jwksResolver.
func newJwksResolver(expireDuration, evictionDuration, refreshInterval time.Duration) *jwksResolver {
	ret := &jwksResolver{
		JwksURICache:     cache.NewTTL(jwksURICacheExpiration, jwksURICacheEviction),
		closing:          make(chan bool, 1),
		expireDuration:   expireDuration,
		evictionDuration: evictionDuration,
		refreshInterval:  refreshInterval,
		httpClient: &http.Client{
			Timeout: jwksHTTPTimeOutInSec * time.Second,

			// TODO: pilot needs to include a collection of root CAs to make external
			// https web request(https://github.com/istio/istio/issues/1419).
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		},
	}

	caCert, err := ioutil.ReadFile(publicRootCABundlePath)
	if err == nil {
		caCertPool := x509.NewCertPool()
		caCertPool.AppendCertsFromPEM(caCert)
		ret.secureHTTPClient = &http.Client{
			Timeout: jwksHTTPTimeOutInSec * time.Second,
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					RootCAs: caCertPool,
				},
			},
		}
	}

	atomic.StoreUint64(&ret.keyChangedCount, 0)
	go ret.refresher()

	return ret
}

// Set jwks_uri through openID discovery if it's not set in auth policy.
func (r *jwksResolver) SetAuthenticationPolicyJwksURIs(policy *authn.Policy) error {
	if policy == nil {
		return fmt.Errorf("invalid nil policy")
	}

	for _, method := range policy.Peers {
		switch method.GetParams().(type) {
		case *authn.PeerAuthenticationMethod_Jwt:
			policyJwt := method.GetJwt()
			if policyJwt.JwksUri == "" {
				uri, err := r.resolveJwksURIUsingOpenID(policyJwt.Issuer)
				if err != nil {
					log.Warnf("Failed to get jwks_uri for issuer %q: %v", policyJwt.Issuer, err)
					return err
				}
				policyJwt.JwksUri = uri
			}
		}
	}
	for _, method := range policy.Origins {
		// JWT is only allowed authentication method type for Origin.
		policyJwt := method.GetJwt()
		if policyJwt.JwksUri == "" {
			uri, err := r.resolveJwksURIUsingOpenID(policyJwt.Issuer)
			if err != nil {
				log.Warnf("Failed to get jwks_uri for issuer %q: %v", policyJwt.Issuer, err)
				return err
			}
			policyJwt.JwksUri = uri
		}
	}

	return nil
}

// GetPublicKey gets JWT public key and cache the key for future use.
func (r *jwksResolver) GetPublicKey(jwksURI string) (string, error) {
	now := time.Now()
	if val, found := r.keyEntries.Load(jwksURI); found {
		e := val.(jwtPubKeyEntry)

		// Return from cache if it's not expired.
		if e.expireTime.After(now) {
			// Update cached key's last used time.
			e.lastUsedTime = now
			r.keyEntries.Store(jwksURI, e)
			return e.pubKey, nil
		}
	}

	// Fetch key if it's not cached, or cached item is expired.
	resp, err := r.getRemoteContent(jwksURI)
	if err != nil {
		log.Errorf("Failed to fetch pubkey from %q: %v", jwksURI, err)
		return "", err
	}

	pubKey := string(resp)
	r.keyEntries.Store(jwksURI, jwtPubKeyEntry{
		pubKey:       pubKey,
		expireTime:   now.Add(r.expireDuration),
		lastUsedTime: now,
	})

	return pubKey, nil
}

// Resolve jwks_uri through openID discovery and cache the jwks_uri for future use.
func (r *jwksResolver) resolveJwksURIUsingOpenID(issuer string) (string, error) {
	// Set policyJwt.JwksUri if the JwksUri could be found in cache.
	if uri, found := r.JwksURICache.Get(issuer); found {
		return uri.(string), nil
	}

	// Try to get jwks_uri through OpenID Discovery.
	body, err := r.getRemoteContent(issuer + openIDDiscoveryCfgURLSuffix)
	if err != nil {
		log.Errorf("Failed to fetch jwks_uri from %q: %v", issuer+openIDDiscoveryCfgURLSuffix, err)
		return "", err
	}
	var data map[string]interface{}
	if err := json.Unmarshal(body, &data); err != nil {
		return "", err
	}

	jwksURI, ok := data["jwks_uri"].(string)
	if !ok {
		return "", fmt.Errorf("invalid jwks_uri %v in openID discovery configuration", data["jwks_uri"])
	}

	// Set JwksUri in cache.
	r.JwksURICache.Set(issuer, jwksURI)

	return jwksURI, nil
}

func (r *jwksResolver) getRemoteContent(uri string) ([]byte, error) {
	u, err := url.Parse(uri)
	if err != nil {
		log.Errorf("Failed to parse %q", uri)
		return nil, err
	}

	client := r.httpClient
	if strings.EqualFold(u.Scheme, "https") {
		// https client may be uninitialized because of root CA bundle missing.
		if r.secureHTTPClient == nil {
			return nil, fmt.Errorf("pilot does not support fetch public key through https endpoint %q", uri)
		}

		client = r.secureHTTPClient
	}

	resp, err := client.Get(uri)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = resp.Body.Close()
	}()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("unsuccessful response from %q", uri)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return body, nil
}

func (r *jwksResolver) refresher() {
	// Wake up once in a while and refresh stale items.
	r.refreshTicker = time.NewTicker(r.refreshInterval)
	for {
		select {
		case now := <-r.refreshTicker.C:
			r.refresh(now)
		case <-r.closing:
			r.refreshTicker.Stop()
			return
		}
	}
}

func (r *jwksResolver) refresh(t time.Time) {
	var wg sync.WaitGroup
	hasChange := false

	r.keyEntries.Range(func(key interface{}, value interface{}) bool {
		now := time.Now()
		jwksURI := key.(string)
		e := value.(jwtPubKeyEntry)

		// Remove cached item if it hasn't been used for a while.
		if now.Sub(e.lastUsedTime) >= r.evictionDuration {
			r.keyEntries.Delete(jwksURI)
			return true
		}

		oldPubKey := e.pubKey

		// key rotation: fetch JWT public key again if it's expired.
		if e.expireTime.Before(t) {
			// Increment the WaitGroup counter.
			wg.Add(1)

			go func() {
				// Decrement the counter when the goroutine completes.
				defer wg.Done()

				resp, err := r.getRemoteContent(jwksURI)
				if err != nil {
					log.Errorf("Cannot fetch JWT public key from %q, %v", jwksURI, err)
					r.keyEntries.Delete(jwksURI)
					return
				}
				newPubKey := string(resp)

				r.keyEntries.Store(jwksURI, jwtPubKeyEntry{
					pubKey:       newPubKey,
					expireTime:   now.Add(r.expireDuration), // Update expireTime even if prev/current keys are the same.
					lastUsedTime: e.lastUsedTime,            // keep original lastUsedTime.
				})

				if oldPubKey != newPubKey {
					hasChange = true
				}
			}()
		}

		return true
	})

	// Wait for all go routine to complete.
	wg.Wait()

	if hasChange {
		atomic.AddUint64(&r.keyChangedCount, 1)
		// TODO(quanlin): send notification to update config and push config to sidecar.
	}
}

// Shut down the refresher job.
// TODO: may need to figure out the right place to call this function.
// (right now calls it from initDiscoveryService in pkg/bootstrap/server.go).
func (r *jwksResolver) Close() {
	r.closing <- true
}
<commit_msg>Allow downstream to change the location of the root ca bundle file (#5798)<commit_after>// Copyright 2018 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package model

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	authn "istio.io/api/authentication/v1alpha1"
	"istio.io/istio/pkg/cache"
	"istio.io/istio/pkg/log"
)

const (
	// https://openid.net/specs/openid-connect-discovery-1_0.html
	// OpenID Providers supporting Discovery MUST make a JSON document available at the path
	// formed by concatenating the string /.well-known/openid-configuration to the Issuer.
	openIDDiscoveryCfgURLSuffix = "/.well-known/openid-configuration"

	// OpenID Discovery web request timeout.
	jwksHTTPTimeOutInSec = 5

	// JwksURI Cache expiration time duration, individual cached JwksURI item will be removed
	// from cache after its duration expires.
	jwksURICacheExpiration = time.Hour * 24

	// JwksURI Cache eviction time duration, cache eviction is done on a periodic basis,
	// jwksURICacheEviction specifies the frequency at which eviction activities take place.
	jwksURICacheEviction = time.Minute * 30

	// JwtPubKeyExpireDuration is the expire duration for JWT public key in the cache.
	// After this duration expire, refresher job will fetch key for the cached item again.
	JwtPubKeyExpireDuration = time.Hour

	// JwtPubKeyEvictionDuration is the life duration for cached item.
	// Cached item will be removed from the cache if it hasn't been used longer than JwtPubKeyEvictionDuration.
	JwtPubKeyEvictionDuration = 24 * 7 * time.Hour

	// JwtPubKeyRefreshInterval is the running interval of JWT pubKey refresh job.
	JwtPubKeyRefreshInterval = time.Minute * 20
)

var (
	// PublicRootCABundlePath is the path of public root CA bundle in pilot container.
	publicRootCABundlePath = "/cacert.pem"
)

// jwtPubKeyEntry is a single cached entry for jwt public key.
type jwtPubKeyEntry struct {
	pubKey string

	// Cached item will be fetched again by refresher job if (time.now >= expireTime).
	expireTime time.Time

	// Cached item's last used time, which is set in GetPublicKey.
	lastUsedTime time.Time
}

// jwksResolver is resolver for jwksURI and jwt public key.
type jwksResolver struct {
	// cache for jwksURI.
	JwksURICache cache.ExpiringCache

	// cache for JWT public key.
	// map key is jwksURI, map value is jwtPubKeyEntry.
	keyEntries sync.Map

	secureHTTPClient *http.Client
	httpClient       *http.Client
	closing          chan bool
	refreshTicker    *time.Ticker

	expireDuration time.Duration

	// Cached key will be removed from cache if (time.now - cachedItem.lastUsedTime >= evictionDuration), this prevents key cache growing indefinitely.
	evictionDuration time.Duration

	// Refresher job running interval.
	refreshInterval time.Duration

	// How may times refresh job has detected JWT public key change happened, used in unit test.
	keyChangedCount uint64
}

// newJwksResolver creates new instance of jwksResolver.
func newJwksResolver(expireDuration, evictionDuration, refreshInterval time.Duration) *jwksResolver {
	ret := &jwksResolver{
		JwksURICache:     cache.NewTTL(jwksURICacheExpiration, jwksURICacheEviction),
		closing:          make(chan bool, 1),
		expireDuration:   expireDuration,
		evictionDuration: evictionDuration,
		refreshInterval:  refreshInterval,
		httpClient: &http.Client{
			Timeout: jwksHTTPTimeOutInSec * time.Second,

			// TODO: pilot needs to include a collection of root CAs to make external
			// https web request(https://github.com/istio/istio/issues/1419).
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		},
	}

	caCert, err := ioutil.ReadFile(publicRootCABundlePath)
	if err == nil {
		caCertPool := x509.NewCertPool()
		caCertPool.AppendCertsFromPEM(caCert)
		ret.secureHTTPClient = &http.Client{
			Timeout: jwksHTTPTimeOutInSec * time.Second,
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					RootCAs: caCertPool,
				},
			},
		}
	}

	atomic.StoreUint64(&ret.keyChangedCount, 0)
	go ret.refresher()

	return ret
}

// Set jwks_uri through openID discovery if it's not set in auth policy.
func (r *jwksResolver) SetAuthenticationPolicyJwksURIs(policy *authn.Policy) error {
	if policy == nil {
		return fmt.Errorf("invalid nil policy")
	}

	for _, method := range policy.Peers {
		switch method.GetParams().(type) {
		case *authn.PeerAuthenticationMethod_Jwt:
			policyJwt := method.GetJwt()
			if policyJwt.JwksUri == "" {
				uri, err := r.resolveJwksURIUsingOpenID(policyJwt.Issuer)
				if err != nil {
					log.Warnf("Failed to get jwks_uri for issuer %q: %v", policyJwt.Issuer, err)
					return err
				}
				policyJwt.JwksUri = uri
			}
		}
	}
	for _, method := range policy.Origins {
		// JWT is only allowed authentication method type for Origin.
		policyJwt := method.GetJwt()
		if policyJwt.JwksUri == "" {
			uri, err := r.resolveJwksURIUsingOpenID(policyJwt.Issuer)
			if err != nil {
				log.Warnf("Failed to get jwks_uri for issuer %q: %v", policyJwt.Issuer, err)
				return err
			}
			policyJwt.JwksUri = uri
		}
	}

	return nil
}

// GetPublicKey gets JWT public key and cache the key for future use.
func (r *jwksResolver) GetPublicKey(jwksURI string) (string, error) {
	now := time.Now()
	if val, found := r.keyEntries.Load(jwksURI); found {
		e := val.(jwtPubKeyEntry)

		// Return from cache if it's not expired.
		if e.expireTime.After(now) {
			// Update cached key's last used time.
			e.lastUsedTime = now
			r.keyEntries.Store(jwksURI, e)
			return e.pubKey, nil
		}
	}

	// Fetch key if it's not cached, or cached item is expired.
	resp, err := r.getRemoteContent(jwksURI)
	if err != nil {
		log.Errorf("Failed to fetch pubkey from %q: %v", jwksURI, err)
		return "", err
	}

	pubKey := string(resp)
	r.keyEntries.Store(jwksURI, jwtPubKeyEntry{
		pubKey:       pubKey,
		expireTime:   now.Add(r.expireDuration),
		lastUsedTime: now,
	})

	return pubKey, nil
}

// Resolve jwks_uri through openID discovery and cache the jwks_uri for future use.
func (r *jwksResolver) resolveJwksURIUsingOpenID(issuer string) (string, error) {
	// Set policyJwt.JwksUri if the JwksUri could be found in cache.
	if uri, found := r.JwksURICache.Get(issuer); found {
		return uri.(string), nil
	}

	// Try to get jwks_uri through OpenID Discovery.
	body, err := r.getRemoteContent(issuer + openIDDiscoveryCfgURLSuffix)
	if err != nil {
		log.Errorf("Failed to fetch jwks_uri from %q: %v", issuer+openIDDiscoveryCfgURLSuffix, err)
		return "", err
	}
	var data map[string]interface{}
	if err := json.Unmarshal(body, &data); err != nil {
		return "", err
	}

	jwksURI, ok := data["jwks_uri"].(string)
	if !ok {
		return "", fmt.Errorf("invalid jwks_uri %v in openID discovery configuration", data["jwks_uri"])
	}

	// Set JwksUri in cache.
	r.JwksURICache.Set(issuer, jwksURI)

	return jwksURI, nil
}

func (r *jwksResolver) getRemoteContent(uri string) ([]byte, error) {
	u, err := url.Parse(uri)
	if err != nil {
		log.Errorf("Failed to parse %q", uri)
		return nil, err
	}

	client := r.httpClient
	if strings.EqualFold(u.Scheme, "https") {
		// https client may be uninitialized because of root CA bundle missing.
		if r.secureHTTPClient == nil {
			return nil, fmt.Errorf("pilot does not support fetch public key through https endpoint %q", uri)
		}

		client = r.secureHTTPClient
	}

	resp, err := client.Get(uri)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = resp.Body.Close()
	}()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("unsuccessful response from %q", uri)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return body, nil
}

func (r *jwksResolver) refresher() {
	// Wake up once in a while and refresh stale items.
	r.refreshTicker = time.NewTicker(r.refreshInterval)
	for {
		select {
		case now := <-r.refreshTicker.C:
			r.refresh(now)
		case <-r.closing:
			r.refreshTicker.Stop()
			return
		}
	}
}

func (r *jwksResolver) refresh(t time.Time) {
	var wg sync.WaitGroup
	hasChange := false

	r.keyEntries.Range(func(key interface{}, value interface{}) bool {
		now := time.Now()
		jwksURI := key.(string)
		e := value.(jwtPubKeyEntry)

		// Remove cached item if it hasn't been used for a while.
		if now.Sub(e.lastUsedTime) >= r.evictionDuration {
			r.keyEntries.Delete(jwksURI)
			return true
		}

		oldPubKey := e.pubKey

		// key rotation: fetch JWT public key again if it's expired.
		if e.expireTime.Before(t) {
			// Increment the WaitGroup counter.
			wg.Add(1)

			go func() {
				// Decrement the counter when the goroutine completes.
				defer wg.Done()

				resp, err := r.getRemoteContent(jwksURI)
				if err != nil {
					log.Errorf("Cannot fetch JWT public key from %q, %v", jwksURI, err)
					r.keyEntries.Delete(jwksURI)
					return
				}
				newPubKey := string(resp)

				r.keyEntries.Store(jwksURI, jwtPubKeyEntry{
					pubKey:       newPubKey,
					expireTime:   now.Add(r.expireDuration), // Update expireTime even if prev/current keys are the same.
					lastUsedTime: e.lastUsedTime,            // keep original lastUsedTime.
				})

				if oldPubKey != newPubKey {
					hasChange = true
				}
			}()
		}

		return true
	})

	// Wait for all go routine to complete.
	wg.Wait()

	if hasChange {
		atomic.AddUint64(&r.keyChangedCount, 1)
		// TODO(quanlin): send notification to update config and push config to sidecar.
	}
}

// Shut down the refresher job.
// TODO: may need to figure out the right place to call this function.
// (right now calls it from initDiscoveryService in pkg/bootstrap/server.go).
func (r *jwksResolver) Close() {
	r.closing <- true
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// 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 conversionaggregator defines the aggregation pipeline for the conversion measurements.
package conversionaggregator

import (
	"encoding/base64"
	"fmt"
	"log"
	"math"
	"reflect"
	"strings"

	"github.com/google/differential-privacy/privacy-on-beam/pbeam"
	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/io/textio"
	"github.com/apache/beam/sdks/go/pkg/beam/transforms/stats"
	"google.golang.org/protobuf/proto"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/conversion"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/secretshare"

	pb "github.com/google/privacy-sandbox-aggregation-service/pipeline/crypto_go_proto"
	// The following packages are required to read files from GCS or local.
	_ "github.com/apache/beam/sdks/go/pkg/beam/io/filesystem/gcs"
	_ "github.com/apache/beam/sdks/go/pkg/beam/io/filesystem/local"
)

// The privacy budget used for adding Laplace noise in the aggregation.
var epsilon = math.Log(3)

const (
	delta    = 1e-5
	minValue = 0
	maxValue = 100
)

func init() {
	beam.RegisterType(reflect.TypeOf((*pb.PartialAggregation)(nil)).Elem())
	beam.RegisterFunction(assemblePartialAggregationFn)
	beam.RegisterFunction(assembleSingleAggregatedReportFn)
	beam.RegisterFunction(extractAggregationKeyValuesFn)
	beam.RegisterFunction(findKeyShareWithMaximumReportIDFn)
	beam.RegisterFunction(formatCompleteAggregationFn)
	beam.RegisterFunction(formatPartialAggregationFn)
	beam.RegisterFunction(mergeAggregationFn)
	beam.RegisterFunction(parsePartialAggregationFn)
}

// PrivacyParams contains parameters required for aggregation with differential privacy.
//
// These parameters are the same with the those defined in
// http://google3/third_party/differential_privacy/go/plume/pbeam/count.go?l=31&rcl=340419070 and
// http://google3/third_party/differential_privacy/go/plume/pbeam/sum.go?l=42&rcl=336106761
type PrivacyParams struct {
	Epsilon, Delta, MinValue, MaxValue float64
}

func extractAggregationKeyValuesFn(data conversion.AggData) (string, int64) {
	return data.AggID, data.ValueShare
}

// aggregateSumDP and aggregateCountDP apply the sum and count aggregations with differential privacy based on the Plume libraries.
//
// Note that in Plume privacy is further protected by occasionally dropping aggregation keys with small counts (http://google3/third_party/differential_privacy/go/library/dpagg/select_partition.go?l=30&rcl=339027652).
func aggregateSumDP(s beam.Scope, pCol pbeam.PrivatePCollection, params PrivacyParams) beam.PCollection {
	s = s.Scope("AggregateSumDP")

	keyToValue := pbeam.ParDo(s, extractAggregationKeyValuesFn, pCol)
	return pbeam.SumPerKey(s, keyToValue, pbeam.SumParams{
		NoiseKind:                pbeam.LaplaceNoise{},
		MaxPartitionsContributed: 1,
		MinValue:                 params.MinValue,
		MaxValue:                 params.MaxValue,
		Epsilon:                  params.Epsilon,
		Delta:                    params.Delta,
	})
}

func extractAggregationKeysFn(data conversion.AggData) string {
	return data.AggID
}

// TODO: We should make sure the count noise is applied deterministically based on a shared seed. For now as we are using the plume count function, the noises may not be consistent for a AggID on different helpers.
func aggregateCountDP(s beam.Scope, pCol pbeam.PrivatePCollection, params PrivacyParams) beam.PCollection {
	s = s.Scope("AggregateCountDP")

	keys := pbeam.ParDo(s, extractAggregationKeysFn, pCol)

	return pbeam.DistinctPrivacyID(s, keys, pbeam.DistinctPrivacyIDParams{
		NoiseKind:                pbeam.LaplaceNoise{},
		MaxPartitionsContributed: 1,
		Epsilon:                  params.Epsilon,
		Delta:                    params.Delta,
	})
}

func assemblePartialAggregationFn(aggID string, countIter, sumIter func(*int64) bool, emit func(string, *pb.PartialAggregation)) {
	var c int64
	countIter(&c)

	var s int64
	sumIter(&s)

	emit(aggID, &pb.PartialAggregation{
		// The int64 summation is directly cast into uint32 intentionally. We are using the overflow behavior of uint32 to handle the summation of secret shares.
		PartialSum:   uint32(s),
		PartialCount: c,
	})
}

func extractKeyFn(key string, value int) string {
	return key
}

func valueToInt64Fn(key string, value int) (string, int64) {
	return key, int64(value)
}

// aggregateOrig counts and sums the share values for each unique AggID in the PCollection<AggData> WITHOUT differential privacy.
func aggregateOrig(s beam.Scope, col beam.PCollection) beam.PCollection {
	cCol := beam.ParDo(s, extractAggregationKeysFn, col)
	aggCountOrig := stats.Count(s, cCol)
	// Convert the value type to be consistent with the DP aggregation results.
	aggCount := beam.ParDo(s, valueToInt64Fn, aggCountOrig)

	sCol := beam.ParDo(s, extractAggregationKeyValuesFn, col)
	aggSum := stats.SumPerKey(s, sCol)

	joined := beam.CoGroupByKey(s, aggCount, aggSum)
	return beam.ParDo(s, assemblePartialAggregationFn, joined)
}

// aggregateDP calculates the count and sum for each unique AggID in the input PCollection<AggData> WITH differential privacy.
//
// Despite adding a Laplace noise in summation, privacy is further protected by dropping AggIDs with small numbers of privacy IDs: http://google3/third_party/differential_privacy/go/library/dpagg/select_partition.go?l=30&rcl=346079936
func aggregateDP(s beam.Scope, col beam.PCollection, params PrivacyParams) beam.PCollection {
	spec := pbeam.NewPrivacySpec(params.Epsilon, params.Delta)
	// There are no privacy identifiers in the conversion data, which is needed by the Plume
	// aggregation functions. We use the ephemeral report ID as the record-level privacy identifier.
	// In the final API, we can achieve user-level bounds by capping contributions on the client side.
	pCol := pbeam.MakePrivateFromStruct(s, col, spec, "ReportID")

	aggCount := aggregateCountDP(s, pCol, PrivacyParams{
		Epsilon: params.Epsilon / 2,
		Delta:   params.Delta / 2,
	})
	aggSum := aggregateSumDP(s, pCol, PrivacyParams{
		Epsilon:  params.Epsilon / 2,
		Delta:    params.Delta / 2,
		MaxValue: params.MaxValue,
		MinValue: params.MinValue,
	})

	joined := beam.CoGroupByKey(s, aggCount, aggSum)
	return beam.ParDo(s, assemblePartialAggregationFn, joined)
}

// Associate the aggregation results with the aggregated key shares.
func assembleSingleAggregatedReportFn(aggID string, keySharesIter func(*string) bool, resultsIter func(**pb.PartialAggregation) bool, emit func(string, *pb.PartialAggregation)) {
	var keyShare string
	keySharesIter(&keyShare)

	var result *pb.PartialAggregation
	resultsIter(&result)
	if result == nil {
		// Emit nothing if there's no aggregation info. Some keys may have been ignored because of differential privacy.
		return
	}

	result.KeyShare = []byte(keyShare)
	emit(aggID, result)
}

// For the same conversion key, we keep the key share with the maximum ephemeral report ID.
func findKeyShareWithMaximumReportIDFn(aggID string, idKeyShareIter func(*conversion.IDKeyShare) bool) (string, string) {
	var minID, aggKeyShare string
	var idKeyShare conversion.IDKeyShare
	for idKeyShareIter(&idKeyShare) {
		if idKeyShare.ReportID > minID {
			aggKeyShare = string(idKeyShare.KeyShare)
			minID = idKeyShare.ReportID
		}
	}
	return aggID, aggKeyShare
}

// AggregateDataShare outputs a PCollection<AggID, PartialAggregation> for the input data, applying differential privacy or not based on the input parameters.
func AggregateDataShare(s beam.Scope, aggIDKeyShare, aggData beam.PCollection, ignorePrivacy bool, params PrivacyParams) beam.PCollection {
	s = s.Scope("AggregateShares")
	aggIDKeys := beam.GroupByKey(s, aggIDKeyShare)
	aggregatedKeyShares := beam.ParDo(s, findKeyShareWithMaximumReportIDFn, aggIDKeys)

	var aggIDResult beam.PCollection
	if ignorePrivacy {
		aggIDResult = aggregateOrig(s, aggData)
	} else {
		aggIDResult = aggregateDP(s, aggData, params)
	}
	joined := beam.CoGroupByKey(s, aggregatedKeyShares, aggIDResult)
	return beam.ParDo(s, assembleSingleAggregatedReportFn, joined)
}

func formatPartialAggregationFn(aggID string, result *pb.PartialAggregation, emit func(string)) error {
	b, err := proto.Marshal(result)
	if err != nil {
		return err
	}
	emit(fmt.Sprintf("%s,%s", base64.StdEncoding.EncodeToString([]byte(aggID)), base64.StdEncoding.EncodeToString(b)))
	return nil
}

// Format and write the partial aggregation results in PCollection<AggID, PartialAggregation>.
func writePartialAggregation(s beam.Scope, aggIDResult beam.PCollection, fileName string) {
	s = s.Scope("WritePartialAggregation")
	formatted := beam.ParDo(s, formatPartialAggregationFn, aggIDResult)
	textio.Write(s, fileName, formatted)
}

// AggregateParams contains necessary parameters for PartialReport aggregation.
type AggregateParams struct {
	PartialReport, ReencryptedKey, PartialAggregation string
	HelperInfo                                        *conversion.ServerPrivateInfo
	IgnorePrivacy                                     bool
}

// AggregatePartialReport calculates the AggID and aggregates data for each unique AggID.
func AggregatePartialReport(scope beam.Scope, aggParams *AggregateParams, privacyParams PrivacyParams) {
	scope = scope.Scope("AggregatePartialReport")

	reportIDKeys := conversion.ReadExponentiatedKeys(scope, aggParams.ReencryptedKey)

	encrypted := conversion.ReadPartialReport(scope, aggParams.PartialReport)
	resharded := beam.Reshuffle(scope, encrypted)
	partialReport := conversion.DecryptPartialReport(scope, resharded, aggParams.HelperInfo.StandardPrivateKey)

	aggIDKeyShare, aggData := conversion.RekeyByAggregationID(scope, reportIDKeys, partialReport, aggParams.HelperInfo.ElGamalPrivateKey, aggParams.HelperInfo.Secret)
	partialAggregation := AggregateDataShare(scope, aggIDKeyShare, aggData, aggParams.IgnorePrivacy, privacyParams)

	writePartialAggregation(scope, partialAggregation, aggParams.PartialAggregation)
}

func parsePartialAggregationFn(line string, emit func(string, *pb.PartialAggregation)) error {
	cols := strings.Split(line, ",")
	if got, want := len(cols), 2; got != want {
		return fmt.Errorf("got %d number of columns in line %q, expected %d", got, line, want)
	}

	bAggID, err := base64.StdEncoding.DecodeString(cols[0])
	if err != nil {
		return err
	}

	bResult, err := base64.StdEncoding.DecodeString(cols[1])
	if err != nil {
		return err
	}

	aggregation := &pb.PartialAggregation{}
	if err := proto.Unmarshal(bResult, aggregation); err != nil {
		return err
	}
	emit(string(bAggID), aggregation)
	return nil
}

func readPartialAggregation(s beam.Scope, partialAggregationFile string) beam.PCollection {
	s = s.Scope("ReadPartialAggregation")
	lines := textio.ReadSdf(s, partialAggregationFile)
	return beam.ParDo(s, parsePartialAggregationFn, lines)
}

// CompleteResult contains the final aggregation result for each conversion key.
type CompleteResult struct {
	ConversionKey string
	Sum           uint32
	Count         int64
}

func mergeAggregationFn(aggID string, pAggIter1 func(**pb.PartialAggregation) bool, pAggIter2 func(**pb.PartialAggregation) bool, emit func(CompleteResult)) error {
	aggregation1 := &pb.PartialAggregation{}
	if !pAggIter1(&aggregation1) {
		log.Printf("expect two shares for aggregation ID %q, missing from helper1", aggID)
		return nil
	}
	aggregation2 := &pb.PartialAggregation{}
	if !pAggIter2(&aggregation2) {
		log.Printf("expect two shares for aggregation ID %q, missing from helper2", aggID)
		return nil
	}
	conversionKey, err := secretshare.CombineByteShares(aggregation1.GetKeyShare(), aggregation2.GetKeyShare())
	if err != nil {
		return err
	}

	emit(CompleteResult{
		ConversionKey: string(conversionKey),
		Sum:           secretshare.CombineIntShares(aggregation1.GetPartialSum(), aggregation2.GetPartialSum()),
		Count:         aggregation1.GetPartialCount() + aggregation2.GetPartialCount(),
	})
	return nil
}

// MergeAggregation combines the partial aggregations to get the complete results.
func MergeAggregation(s beam.Scope, pAgg1, pAgg2 beam.PCollection) beam.PCollection {
	s = s.Scope("MergePartialAggregations")
	joined := beam.CoGroupByKey(s, pAgg1, pAgg2)
	return beam.ParDo(s, mergeAggregationFn, joined)
}

func formatCompleteAggregationFn(result CompleteResult) string {
	return fmt.Sprintf("%s,%d,%d", string(result.ConversionKey), result.Count, result.Sum)
}

func writeCompleteAggregation(s beam.Scope, aggIDResult beam.PCollection, fileName string) {
	s = s.Scope("WriteCompleteAggregation")
	formatted := beam.ParDo(s, formatCompleteAggregationFn, aggIDResult)
	textio.Write(s, fileName, formatted)
}

// MergePartialAggregation calculates the complete result from the aggregated shares.
func MergePartialAggregation(scope beam.Scope, partialAggregationFile1, partialAggregationFile2, completeAggregationFile string) {
	scope = scope.Scope("CompleteAggregate")

	partialAggregation1 := readPartialAggregation(scope, partialAggregationFile1)
	partialAggregation2 := readPartialAggregation(scope, partialAggregationFile2)
	completeAggregation := MergeAggregation(scope, partialAggregation1, partialAggregation2)
	writeCompleteAggregation(scope, completeAggregation, completeAggregationFile)
}
<commit_msg>Internal change<commit_after>// Copyright 2021 Google LLC
//
// 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 conversionaggregator defines the aggregation pipeline for the conversion measurements.
package conversionaggregator

import (
	"encoding/base64"
	"fmt"
	"log"
	"math"
	"reflect"
	"strings"

	"github.com/google/differential-privacy/privacy-on-beam/pbeam"
	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/io/textio"
	"github.com/apache/beam/sdks/go/pkg/beam/transforms/stats"
	"google.golang.org/protobuf/proto"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/conversion"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/secretshare"

	pb "github.com/google/privacy-sandbox-aggregation-service/pipeline/crypto_go_proto"
	// The following packages are required to read files from GCS or local.
	_ "github.com/apache/beam/sdks/go/pkg/beam/io/filesystem/gcs"
	_ "github.com/apache/beam/sdks/go/pkg/beam/io/filesystem/local"
)

// The privacy budget used for adding Laplace noise in the aggregation.
var epsilon = math.Log(3)

const (
	delta    = 1e-5
	minValue = 0
	maxValue = 100
)

func init() {
	beam.RegisterType(reflect.TypeOf((*pb.PartialAggregation)(nil)).Elem())
	beam.RegisterFunction(assemblePartialAggregationFn)
	beam.RegisterFunction(assembleSingleAggregatedReportFn)
	beam.RegisterFunction(extractAggregationKeyValuesFn)
	beam.RegisterFunction(findKeyShareWithMaximumReportIDFn)
	beam.RegisterFunction(formatCompleteAggregationFn)
	beam.RegisterFunction(formatPartialAggregationFn)
	beam.RegisterFunction(mergeAggregationFn)
	beam.RegisterFunction(parsePartialAggregationFn)
}

// PrivacyParams contains parameters required for aggregation with differential privacy.
//
// These parameters are the same with the those defined in
// http://google3/third_party/differential_privacy/go/plume/pbeam/count.go?l=31&rcl=340419070 and
// http://google3/third_party/differential_privacy/go/plume/pbeam/sum.go?l=42&rcl=336106761
type PrivacyParams struct {
	Epsilon, Delta, MinValue, MaxValue float64
}

func extractAggregationKeyValuesFn(data conversion.AggData) (string, int64) {
	return data.AggID, data.ValueShare
}

// aggregateSumDP and aggregateCountDP apply the sum and count aggregations with differential privacy based on the Plume libraries.
//
// Note that in Plume privacy is further protected by occasionally dropping aggregation keys with small counts (http://google3/third_party/differential_privacy/go/library/dpagg/select_partition.go?l=30&rcl=339027652).
func aggregateSumDP(s beam.Scope, pCol pbeam.PrivatePCollection, params PrivacyParams) beam.PCollection {
	s = s.Scope("AggregateSumDP")

	keyToValue := pbeam.ParDo(s, extractAggregationKeyValuesFn, pCol)
	return pbeam.SumPerKey(s, keyToValue, pbeam.SumParams{
		NoiseKind:                pbeam.LaplaceNoise{},
		MaxPartitionsContributed: 1,
		MinValue:                 params.MinValue,
		MaxValue:                 params.MaxValue,
		Epsilon:                  params.Epsilon,
		Delta:                    params.Delta,
	})
}

func extractAggregationKeysFn(data conversion.AggData) string {
	return data.AggID
}

// TODO: We should make sure the count noise is applied deterministically based on a shared seed. For now as we are using the plume count function, the noises may not be consistent for a AggID on different helpers.
func aggregateCountDP(s beam.Scope, pCol pbeam.PrivatePCollection, params PrivacyParams) beam.PCollection {
	s = s.Scope("AggregateCountDP")

	keys := pbeam.ParDo(s, extractAggregationKeysFn, pCol)

	return pbeam.DistinctPrivacyID(s, keys, pbeam.DistinctPrivacyIDParams{
		NoiseKind:                pbeam.LaplaceNoise{},
		MaxPartitionsContributed: 1,
		Epsilon:                  params.Epsilon,
		Delta:                    params.Delta,
	})
}

func assemblePartialAggregationFn(aggID string, countIter, sumIter func(*int64) bool, emit func(string, *pb.PartialAggregation)) {
	var c int64
	countIter(&c)

	var s int64
	sumIter(&s)

	emit(aggID, &pb.PartialAggregation{
		// The int64 summation is directly cast into uint32 intentionally. We are using the overflow behavior of uint32 to handle the summation of secret shares.
		PartialSum:   uint32(s),
		PartialCount: c,
	})
}

func extractKeyFn(key string, value int) string {
	return key
}

func valueToInt64Fn(key string, value int) (string, int64) {
	return key, int64(value)
}

// aggregateOrig counts and sums the share values for each unique AggID in the PCollection<AggData> WITHOUT differential privacy.
func aggregateOrig(s beam.Scope, col beam.PCollection) beam.PCollection {
	cCol := beam.ParDo(s, extractAggregationKeysFn, col)
	aggCountOrig := stats.Count(s, cCol)
	// Convert the value type to be consistent with the DP aggregation results.
	aggCount := beam.ParDo(s, valueToInt64Fn, aggCountOrig)

	sCol := beam.ParDo(s, extractAggregationKeyValuesFn, col)
	aggSum := stats.SumPerKey(s, sCol)

	joined := beam.CoGroupByKey(s, aggCount, aggSum)
	return beam.ParDo(s, assemblePartialAggregationFn, joined)
}

// aggregateDP calculates the count and sum for each unique AggID in the input PCollection<AggData> WITH differential privacy.
//
// Despite adding a Laplace noise in summation, privacy is further protected by dropping AggIDs with small numbers of privacy IDs: http://google3/third_party/differential_privacy/go/library/dpagg/select_partition.go?l=30&rcl=346079936
func aggregateDP(s beam.Scope, col beam.PCollection, params PrivacyParams) beam.PCollection {
	spec := pbeam.NewPrivacySpec(params.Epsilon, params.Delta)
	// There are no privacy identifiers in the conversion data, which is needed by the Plume
	// aggregation functions. We use the ephemeral report ID as the record-level privacy identifier.
	// In the final API, we can achieve user-level bounds by capping contributions on the client side.
	pCol := pbeam.MakePrivateFromStruct(s, col, spec, "ReportID")

	aggCount := aggregateCountDP(s, pCol, PrivacyParams{
		Epsilon: params.Epsilon / 2,
		Delta:   params.Delta / 2,
	})
	aggSum := aggregateSumDP(s, pCol, PrivacyParams{
		Epsilon:  params.Epsilon / 2,
		Delta:    params.Delta / 2,
		MaxValue: params.MaxValue,
		MinValue: params.MinValue,
	})

	joined := beam.CoGroupByKey(s, aggCount, aggSum)
	return beam.ParDo(s, assemblePartialAggregationFn, joined)
}

// Associate the aggregation results with the aggregated key shares.
func assembleSingleAggregatedReportFn(aggID string, keySharesIter func(*string) bool, resultsIter func(**pb.PartialAggregation) bool, emit func(string, *pb.PartialAggregation)) {
	var keyShare string
	keySharesIter(&keyShare)

	var result *pb.PartialAggregation
	resultsIter(&result)
	if result == nil {
		// Emit nothing if there's no aggregation info. Some keys may have been ignored because of differential privacy.
		return
	}

	result.KeyShare = []byte(keyShare)
	emit(aggID, result)
}

// For the same conversion key, we keep the key share with the maximum ephemeral report ID.
func findKeyShareWithMaximumReportIDFn(aggID string, idKeyShareIter func(*conversion.IDKeyShare) bool) (string, string) {
	var maxID, aggKeyShare string
	var idKeyShare conversion.IDKeyShare
	for idKeyShareIter(&idKeyShare) {
		if idKeyShare.ReportID > maxID {
			aggKeyShare = string(idKeyShare.KeyShare)
			maxID = idKeyShare.ReportID
		}
	}
	return aggID, aggKeyShare
}

// AggregateDataShare outputs a PCollection<AggID, PartialAggregation> for the input data, applying differential privacy or not based on the input parameters.
func AggregateDataShare(s beam.Scope, aggIDKeyShare, aggData beam.PCollection, ignorePrivacy bool, params PrivacyParams) beam.PCollection {
	s = s.Scope("AggregateShares")
	aggIDKeys := beam.GroupByKey(s, aggIDKeyShare)
	aggregatedKeyShares := beam.ParDo(s, findKeyShareWithMaximumReportIDFn, aggIDKeys)

	var aggIDResult beam.PCollection
	if ignorePrivacy {
		aggIDResult = aggregateOrig(s, aggData)
	} else {
		aggIDResult = aggregateDP(s, aggData, params)
	}
	joined := beam.CoGroupByKey(s, aggregatedKeyShares, aggIDResult)
	return beam.ParDo(s, assembleSingleAggregatedReportFn, joined)
}

func formatPartialAggregationFn(aggID string, result *pb.PartialAggregation, emit func(string)) error {
	b, err := proto.Marshal(result)
	if err != nil {
		return err
	}
	emit(fmt.Sprintf("%s,%s", base64.StdEncoding.EncodeToString([]byte(aggID)), base64.StdEncoding.EncodeToString(b)))
	return nil
}

// Format and write the partial aggregation results in PCollection<AggID, PartialAggregation>.
func writePartialAggregation(s beam.Scope, aggIDResult beam.PCollection, fileName string) {
	s = s.Scope("WritePartialAggregation")
	formatted := beam.ParDo(s, formatPartialAggregationFn, aggIDResult)
	textio.Write(s, fileName, formatted)
}

// AggregateParams contains necessary parameters for PartialReport aggregation.
type AggregateParams struct {
	PartialReport, ReencryptedKey, PartialAggregation string
	HelperInfo                                        *conversion.ServerPrivateInfo
	IgnorePrivacy                                     bool
}

// AggregatePartialReport calculates the AggID and aggregates data for each unique AggID.
func AggregatePartialReport(scope beam.Scope, aggParams *AggregateParams, privacyParams PrivacyParams) {
	scope = scope.Scope("AggregatePartialReport")

	reportIDKeys := conversion.ReadExponentiatedKeys(scope, aggParams.ReencryptedKey)

	encrypted := conversion.ReadPartialReport(scope, aggParams.PartialReport)
	resharded := beam.Reshuffle(scope, encrypted)
	partialReport := conversion.DecryptPartialReport(scope, resharded, aggParams.HelperInfo.StandardPrivateKey)

	aggIDKeyShare, aggData := conversion.RekeyByAggregationID(scope, reportIDKeys, partialReport, aggParams.HelperInfo.ElGamalPrivateKey, aggParams.HelperInfo.Secret)
	partialAggregation := AggregateDataShare(scope, aggIDKeyShare, aggData, aggParams.IgnorePrivacy, privacyParams)

	writePartialAggregation(scope, partialAggregation, aggParams.PartialAggregation)
}

func parsePartialAggregationFn(line string, emit func(string, *pb.PartialAggregation)) error {
	cols := strings.Split(line, ",")
	if got, want := len(cols), 2; got != want {
		return fmt.Errorf("got %d number of columns in line %q, expected %d", got, line, want)
	}

	bAggID, err := base64.StdEncoding.DecodeString(cols[0])
	if err != nil {
		return err
	}

	bResult, err := base64.StdEncoding.DecodeString(cols[1])
	if err != nil {
		return err
	}

	aggregation := &pb.PartialAggregation{}
	if err := proto.Unmarshal(bResult, aggregation); err != nil {
		return err
	}
	emit(string(bAggID), aggregation)
	return nil
}

func readPartialAggregation(s beam.Scope, partialAggregationFile string) beam.PCollection {
	s = s.Scope("ReadPartialAggregation")
	lines := textio.ReadSdf(s, partialAggregationFile)
	return beam.ParDo(s, parsePartialAggregationFn, lines)
}

// CompleteResult contains the final aggregation result for each conversion key.
type CompleteResult struct {
	ConversionKey string
	Sum           uint32
	Count         int64
}

func mergeAggregationFn(aggID string, pAggIter1 func(**pb.PartialAggregation) bool, pAggIter2 func(**pb.PartialAggregation) bool, emit func(CompleteResult)) error {
	aggregation1 := &pb.PartialAggregation{}
	if !pAggIter1(&aggregation1) {
		log.Printf("expect two shares for aggregation ID %q, missing from helper1", aggID)
		return nil
	}
	aggregation2 := &pb.PartialAggregation{}
	if !pAggIter2(&aggregation2) {
		log.Printf("expect two shares for aggregation ID %q, missing from helper2", aggID)
		return nil
	}
	conversionKey, err := secretshare.CombineByteShares(aggregation1.GetKeyShare(), aggregation2.GetKeyShare())
	if err != nil {
		return err
	}

	emit(CompleteResult{
		ConversionKey: string(conversionKey),
		Sum:           secretshare.CombineIntShares(aggregation1.GetPartialSum(), aggregation2.GetPartialSum()),
		Count:         aggregation1.GetPartialCount() + aggregation2.GetPartialCount(),
	})
	return nil
}

// MergeAggregation combines the partial aggregations to get the complete results.
func MergeAggregation(s beam.Scope, pAgg1, pAgg2 beam.PCollection) beam.PCollection {
	s = s.Scope("MergePartialAggregations")
	joined := beam.CoGroupByKey(s, pAgg1, pAgg2)
	return beam.ParDo(s, mergeAggregationFn, joined)
}

func formatCompleteAggregationFn(result CompleteResult) string {
	return fmt.Sprintf("%s,%d,%d", string(result.ConversionKey), result.Count, result.Sum)
}

func writeCompleteAggregation(s beam.Scope, aggIDResult beam.PCollection, fileName string) {
	s = s.Scope("WriteCompleteAggregation")
	formatted := beam.ParDo(s, formatCompleteAggregationFn, aggIDResult)
	textio.Write(s, fileName, formatted)
}

// MergePartialAggregation calculates the complete result from the aggregated shares.
func MergePartialAggregation(scope beam.Scope, partialAggregationFile1, partialAggregationFile2, completeAggregationFile string) {
	scope = scope.Scope("CompleteAggregate")

	partialAggregation1 := readPartialAggregation(scope, partialAggregationFile1)
	partialAggregation2 := readPartialAggregation(scope, partialAggregationFile2)
	completeAggregation := MergeAggregation(scope, partialAggregation1, partialAggregation2)
	writeCompleteAggregation(scope, completeAggregation, completeAggregationFile)
}
<|endoftext|>


<commit_before>package httprecorder

import (
	"bytes"
	"io"
	"net/http"
	"strings"
	"sync"
)

type Recorder struct {
	Client            *http.Client
	originalTransport http.RoundTripper
}

type CompiledRequest struct {
	Request struct {
		Method  string            `json:"method,omitempty"`
		URL     string            `json:"url,omitempty"`
		Headers map[string]string `json:"headers,omitempty"`
		Body    string            `json:"body,omitempty"`
	} `json:"request,omitempty"`
	Response struct {
		Headers map[string]string `json:"headers,omitempty"`
		Body    string            `json:"body,omitempty"`
	} `json:"response,omitempty"`
}

func NewWithClient(client *http.Client) *Recorder {
	r := &Recorder{Client: client}
	r.originalTransport = client.Transport
	client.Transport = &roundTripRecorder{RoundTripper: r.originalTransport}
	return r
}

func (r *Recorder) ResetClient() {
	r.Client.Transport = r.originalTransport
}

func (r *Recorder) GetRequests() []*CompiledRequest {
	t := r.Client.Transport.(*roundTripRecorder)
	reqs := t.requests
	t.requests = t.requests[:0]
	compiledReqs := make([]*CompiledRequest, len(reqs))
	for i, r := range reqs {
		compiledReqs[i] = compileRequest(r)
	}
	return compiledReqs
}

type request struct {
	req     *http.Request
	res     *http.Response
	reqBody *bytes.Buffer
	resBody *bytes.Buffer
}

type roundTripRecorder struct {
	RoundTripper http.RoundTripper
	requests     []*request
	mtx          sync.Mutex
}

func (r *roundTripRecorder) RoundTrip(req *http.Request) (*http.Response, error) {
	reqBuf, resBuf := &bytes.Buffer{}, &bytes.Buffer{}
	if req.Body != nil {
		req.Body = readCloser{req.Body, io.TeeReader(req.Body, reqBuf)}
	}
	res, err := r.RoundTripper.RoundTrip(req)
	if err != nil {
		return nil, err
	}
	log := &request{req, res, reqBuf, resBuf}
	res.Body = readCloser{res.Body, io.TeeReader(res.Body, resBuf)}
	r.mtx.Lock()
	r.requests = append(r.requests, log)
	r.mtx.Unlock()
	return res, nil
}

type readCloser struct {
	io.Closer
	io.Reader
}

var excludeHeaders = map[string]struct{}{
	"Host":              {},
	"Content-Length":    {},
	"Transfer-Encoding": {},
	"Trailer":           {},
	"User-Agent":        {},
	"Date":              {},
}

func compileRequest(r *request) *CompiledRequest {
	res := &CompiledRequest{}

	// request
	res.Request.Method = r.req.Method

	uri := r.req.URL
	uriStr := uri.Path
	if uri.RawQuery != "" {
		uriStr = uriStr + "?" + uri.RawQuery
	}
	if uri.Fragment != "" {
		uriStr = uriStr + "#" + uri.Fragment
	}
	res.Request.URL = uriStr

	res.Request.Headers = make(map[string]string, len(r.req.Header))
	for k, values := range r.req.Header {
		if _, ok := excludeHeaders[k]; ok {
			continue
		}
		res.Request.Headers[k] = strings.Join(values, ",")
	}

	if r.req.Body != nil {
		res.Request.Body = r.reqBody.String()
	}

	// response
	res.Response.Headers = make(map[string]string, len(r.res.Header))
	for k, values := range r.res.Header {
		if _, ok := excludeHeaders[k]; ok {
			continue
		}
		res.Response.Headers[k] = strings.Join(values, ",")
	}

	resBody := r.resBody.String()
	if len(resBody) > 0 {
		res.Response.Body = resBody
	}

	return res
}
<commit_msg>pkg/httprecorder: Don't require http.Client.Transport to be set<commit_after>package httprecorder

import (
	"bytes"
	"io"
	"net/http"
	"strings"
	"sync"
)

type Recorder struct {
	Client            *http.Client
	originalTransport http.RoundTripper
}

type CompiledRequest struct {
	Request struct {
		Method  string            `json:"method,omitempty"`
		URL     string            `json:"url,omitempty"`
		Headers map[string]string `json:"headers,omitempty"`
		Body    string            `json:"body,omitempty"`
	} `json:"request,omitempty"`
	Response struct {
		Headers map[string]string `json:"headers,omitempty"`
		Body    string            `json:"body,omitempty"`
	} `json:"response,omitempty"`
}

func NewWithClient(client *http.Client) *Recorder {
	r := &Recorder{Client: client}
	r.originalTransport = client.Transport
	if r.originalTransport == nil {
		r.originalTransport = &http.Transport{}
	}
	client.Transport = &roundTripRecorder{RoundTripper: r.originalTransport}
	return r
}

func (r *Recorder) ResetClient() {
	r.Client.Transport = r.originalTransport
}

func (r *Recorder) GetRequests() []*CompiledRequest {
	t := r.Client.Transport.(*roundTripRecorder)
	reqs := t.requests
	t.requests = t.requests[:0]
	compiledReqs := make([]*CompiledRequest, len(reqs))
	for i, r := range reqs {
		compiledReqs[i] = compileRequest(r)
	}
	return compiledReqs
}

type request struct {
	req     *http.Request
	res     *http.Response
	reqBody *bytes.Buffer
	resBody *bytes.Buffer
}

type roundTripRecorder struct {
	RoundTripper http.RoundTripper
	requests     []*request
	mtx          sync.Mutex
}

func (r *roundTripRecorder) RoundTrip(req *http.Request) (*http.Response, error) {
	reqBuf, resBuf := &bytes.Buffer{}, &bytes.Buffer{}
	if req.Body != nil {
		req.Body = readCloser{req.Body, io.TeeReader(req.Body, reqBuf)}
	}
	res, err := r.RoundTripper.RoundTrip(req)
	if err != nil {
		return nil, err
	}
	log := &request{req, res, reqBuf, resBuf}
	res.Body = readCloser{res.Body, io.TeeReader(res.Body, resBuf)}
	r.mtx.Lock()
	r.requests = append(r.requests, log)
	r.mtx.Unlock()
	return res, nil
}

type readCloser struct {
	io.Closer
	io.Reader
}

var excludeHeaders = map[string]struct{}{
	"Host":              {},
	"Content-Length":    {},
	"Transfer-Encoding": {},
	"Trailer":           {},
	"User-Agent":        {},
	"Date":              {},
}

func compileRequest(r *request) *CompiledRequest {
	res := &CompiledRequest{}

	// request
	res.Request.Method = r.req.Method

	uri := r.req.URL
	uriStr := uri.Path
	if uri.RawQuery != "" {
		uriStr = uriStr + "?" + uri.RawQuery
	}
	if uri.Fragment != "" {
		uriStr = uriStr + "#" + uri.Fragment
	}
	res.Request.URL = uriStr

	res.Request.Headers = make(map[string]string, len(r.req.Header))
	for k, values := range r.req.Header {
		if _, ok := excludeHeaders[k]; ok {
			continue
		}
		res.Request.Headers[k] = strings.Join(values, ",")
	}

	if r.req.Body != nil {
		res.Request.Body = r.reqBody.String()
	}

	// response
	res.Response.Headers = make(map[string]string, len(r.res.Header))
	for k, values := range r.res.Header {
		if _, ok := excludeHeaders[k]; ok {
			continue
		}
		res.Response.Headers[k] = strings.Join(values, ",")
	}

	resBody := r.resBody.String()
	if len(resBody) > 0 {
		res.Response.Body = resBody
	}

	return res
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors.

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

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

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

package ingress

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

	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/util/sets"
	corev1listers "k8s.io/client-go/listers/core/v1"
	"knative.dev/net-kourier/pkg/config"
	"knative.dev/networking/pkg/apis/networking/v1alpha1"
	"knative.dev/networking/pkg/status"
)

func NewProbeTargetLister(logger *zap.SugaredLogger, endpointsLister corev1listers.EndpointsLister) status.ProbeTargetLister {
	return &gatewayPodTargetLister{
		logger:          logger,
		endpointsLister: endpointsLister,
	}
}

type gatewayPodTargetLister struct {
	logger          *zap.SugaredLogger
	endpointsLister corev1listers.EndpointsLister
}

func (l *gatewayPodTargetLister) ListProbeTargets(ctx context.Context, ing *v1alpha1.Ingress) ([]status.ProbeTarget, error) {
	eps, err := l.endpointsLister.Endpoints(config.GatewayNamespace()).Get(config.InternalServiceName)
	if err != nil {
		return nil, fmt.Errorf("failed to get internal service: %w", err)
	}

	var readyIPs []string
	for _, sub := range eps.Subsets {
		for _, address := range sub.Addresses {
			readyIPs = append(readyIPs, address.IP)
		}
	}
	if len(readyIPs) == 0 {
		return nil, fmt.Errorf("no gateway pods available")
	}
	return l.getIngressUrls(ing, readyIPs)
}

func (l *gatewayPodTargetLister) getIngressUrls(ing *v1alpha1.Ingress, gatewayIps []string) ([]status.ProbeTarget, error) {
	ips := sets.NewString(gatewayIps...)

	targets := make([]status.ProbeTarget, 0, len(ing.Spec.Rules))
	for _, rule := range ing.Spec.Rules {
		var target status.ProbeTarget

		domains := rule.Hosts
		scheme := "http"

		if rule.Visibility == v1alpha1.IngressVisibilityExternalIP {
			target = status.ProbeTarget{
				PodIPs: ips,
			}
			if len(ing.Spec.TLS) != 0 {
				target.PodPort = strconv.Itoa(int(config.HTTPSPortExternal))
				target.URLs = domainsToURL(domains, "https")
			} else {
				target.PodPort = strconv.Itoa(int(config.HTTPPortExternal))
				target.URLs = domainsToURL(domains, scheme)
			}
			targets = append(targets, target)
		} else {
			target = status.ProbeTarget{
				PodIPs:  ips,
				PodPort: strconv.Itoa(int(config.HTTPPortInternal)),
				URLs:    domainsToURL(domains, scheme),
			}
		}

		targets = append(targets, target)

	}
	return targets, nil
}

func domainsToURL(domains []string, scheme string) []*url.URL {
	urls := make([]*url.URL, 0, len(domains))
	for _, domain := range domains {
		url := &url.URL{
			Scheme: scheme,
			Host:   domain,
			Path:   "/",
		}
		urls = append(urls, url)
	}
	return urls
}
<commit_msg>Do not add targets twice (#584)<commit_after>/*
Copyright 2020 The Knative Authors.

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

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

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

package ingress

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

	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/util/sets"
	corev1listers "k8s.io/client-go/listers/core/v1"
	"knative.dev/net-kourier/pkg/config"
	"knative.dev/networking/pkg/apis/networking/v1alpha1"
	"knative.dev/networking/pkg/status"
)

func NewProbeTargetLister(logger *zap.SugaredLogger, endpointsLister corev1listers.EndpointsLister) status.ProbeTargetLister {
	return &gatewayPodTargetLister{
		logger:          logger,
		endpointsLister: endpointsLister,
	}
}

type gatewayPodTargetLister struct {
	logger          *zap.SugaredLogger
	endpointsLister corev1listers.EndpointsLister
}

func (l *gatewayPodTargetLister) ListProbeTargets(ctx context.Context, ing *v1alpha1.Ingress) ([]status.ProbeTarget, error) {
	eps, err := l.endpointsLister.Endpoints(config.GatewayNamespace()).Get(config.InternalServiceName)
	if err != nil {
		return nil, fmt.Errorf("failed to get internal service: %w", err)
	}

	var readyIPs []string
	for _, sub := range eps.Subsets {
		for _, address := range sub.Addresses {
			readyIPs = append(readyIPs, address.IP)
		}
	}
	if len(readyIPs) == 0 {
		return nil, fmt.Errorf("no gateway pods available")
	}
	return l.getIngressUrls(ing, readyIPs)
}

func (l *gatewayPodTargetLister) getIngressUrls(ing *v1alpha1.Ingress, gatewayIps []string) ([]status.ProbeTarget, error) {
	ips := sets.NewString(gatewayIps...)

	targets := make([]status.ProbeTarget, 0, len(ing.Spec.Rules))
	for _, rule := range ing.Spec.Rules {
		var target status.ProbeTarget

		domains := rule.Hosts
		scheme := "http"

		if rule.Visibility == v1alpha1.IngressVisibilityExternalIP {
			target = status.ProbeTarget{
				PodIPs: ips,
			}
			if len(ing.Spec.TLS) != 0 {
				target.PodPort = strconv.Itoa(int(config.HTTPSPortExternal))
				target.URLs = domainsToURL(domains, "https")
			} else {
				target.PodPort = strconv.Itoa(int(config.HTTPPortExternal))
				target.URLs = domainsToURL(domains, scheme)
			}
		} else {
			target = status.ProbeTarget{
				PodIPs:  ips,
				PodPort: strconv.Itoa(int(config.HTTPPortInternal)),
				URLs:    domainsToURL(domains, scheme),
			}
		}

		targets = append(targets, target)

	}
	return targets, nil
}

func domainsToURL(domains []string, scheme string) []*url.URL {
	urls := make([]*url.URL, 0, len(domains))
	for _, domain := range domains {
		url := &url.URL{
			Scheme: scheme,
			Host:   domain,
			Path:   "/",
		}
		urls = append(urls, url)
	}
	return urls
}
<|endoftext|>