<commit_before>package gce

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/emccode/rexray/core"
	"github.com/emccode/rexray/core/config"
	"github.com/emccode/rexray/core/errors"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/compute/v1"
	"io/ioutil"
	"regexp"
	"strconv"
	"strings"
)

const providerName = "gce"

// The GCE storage driver.
type driver struct {
	client  *compute.Service
	r       *core.RexRay
	zone    string
	project string
}

func ef() errors.Fields {
	return errors.Fields{
		"provider": providerName,
	}
}

func eff(fields errors.Fields) map[string]interface{} {
	errFields := map[string]interface{}{
		"provider": providerName,
	}
	if fields != nil {
		for k, v := range fields {
			errFields[k] = v
		}
	}
	return errFields
}

func init() {
	core.RegisterDriver(providerName, newDriver)
	config.Register(configRegistration())
}

func newDriver() core.Driver {
	return &driver{}
}

func (d *driver) Init(r *core.RexRay) error {
	d.r = r

	var err error

	d.zone = d.r.Config.GetString("gce.zone")
	d.project = d.r.Config.GetString("gce.project")
	serviceAccountJSON, err := ioutil.ReadFile(d.r.Config.GetString("gce.keyfile"))
	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not read service account credentials file, %s => {%s}", d.r.Config.GetString("gce.keyfile"), err)
	}

	config, err := google.JWTConfigFromJSON(serviceAccountJSON,
		compute.ComputeScope,
	)
	client, err := compute.New(config.Client(context.Background()))

	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not create compute client => {%s}", err)
	}
	d.client = client
	log.WithField("provider", providerName).Info("storage driver initialized")

	return nil
}

func (d *driver) Name() string {
	return providerName
}

func (d *driver) GetVolumeMapping() ([]*core.BlockDevice, error) {
	log.WithField("provider", providerName).Debug("GetVolumeMapping")
	return nil, nil
}

func (d *driver) GetInstance() (*core.Instance, error) {
	log.WithField("provider", providerName).Debug("GetInstance")
	return nil, nil
}

func (d *driver) CreateSnapshot(
	runAsync bool,
	snapshotName, volumeID, description string) ([]*core.Snapshot, error) {

	log.WithField("provider", providerName).Debug("CreateSnapshot")
	return nil, nil

}

func (d *driver) GetSnapshot(
	volumeID, snapshotID, snapshotName string) ([]*core.Snapshot, error) {

	log.WithField("provider", providerName).Debug("GetSnapshot")
	return nil, nil
}

func (d *driver) RemoveSnapshot(snapshotID string) error {
	log.WithField("provider", providerName).Debug("RemoveSnapshot")
	return nil
}

func (d *driver) GetDeviceNextAvailable() (string, error) {
	letters := []string{
		"a", "b", "c", "d", "e", "f", "g", "h",
		"i", "j", "k", "l", "m", "n", "o", "p"}

	blockDeviceNames := make(map[string]bool)

	blockDeviceMapping, err := d.GetVolumeMapping()
	if err != nil {
		return "", err
	}

	for _, blockDevice := range blockDeviceMapping {
		re, _ := regexp.Compile(`^/dev/xvd([a-z])`)
		res := re.FindStringSubmatch(blockDevice.DeviceName)
		if len(res) > 0 {
			blockDeviceNames[res[1]] = true
		}
	}

	localDevices, err := getLocalDevices()
	if err != nil {
		return "", err
	}

	for _, localDevice := range localDevices {
		re, _ := regexp.Compile(`^xvd([a-z])`)
		res := re.FindStringSubmatch(localDevice)
		if len(res) > 0 {
			blockDeviceNames[res[1]] = true
		}
	}

	for _, letter := range letters {
		if !blockDeviceNames[letter] {
			nextDeviceName := "/dev/xvd" + letter
			log.Println("Got next device name: " + nextDeviceName)
			return nextDeviceName, nil
		}
	}
	return "", errors.New("No available device")
}

func getLocalDevices() (deviceNames []string, err error) {
	file := "/proc/partitions"
	contentBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return []string{}, err
	}

	content := string(contentBytes)

	lines := strings.Split(content, "\n")
	for _, line := range lines[2:] {
		fields := strings.Fields(line)
		if len(fields) == 4 {
			deviceNames = append(deviceNames, fields[3])
		}
	}

	return deviceNames, nil
}

func (d *driver) CreateVolume(
	runAsync bool, volumeName, volumeID, snapshotID, volumeType string,
	IOPS, size int64, availabilityZone string) (*core.Volume, error) {
	log.WithField("provider", providerName).Debug("CreateVolume")
	return nil, nil

}

func (d *driver) createVolumeCreateSnapshot(
	volumeID string, snapshotID string) (string, error) {
	log.WithField("provider", providerName).Debug("CreateVolumeCreateSnapshot")
	return "", nil

}

func (d *driver) GetVolume(
	volumeID, volumeName string) ([]*core.Volume, error) {
	log.WithField("provider", providerName).Debugf("GetVolume :%s %s", volumeID, volumeName)

	query := d.client.Disks.List(d.project, d.zone)
	if volumeID != "" {
		query.Filter(fmt.Sprintf("id eq %s", volumeID))
	}
	if volumeName != "" {
		query.Filter(fmt.Sprintf("name eq %s", volumeName))
	}
	var attachments []*core.VolumeAttachment
	instances, err := d.client.Instances.List(d.project, d.zone).Do()
	for _, instance := range instances.Items {
		for _, disk := range instance.Disks {
			attachment := &core.VolumeAttachment{
				InstanceID: strconv.FormatUint(instance.Id, 10),
				DeviceName: disk.DeviceName,
				Status:     disk.Mode,
				VolumeID:   disk.Source,
			}
			attachments = append(attachments, attachment)

		}
	}

	disks, err := query.Do()
	if err != nil {
		return []*core.Volume{}, err
	}
	var volumesSD []*core.Volume
	for _, disk := range disks.Items {
		var diskAttachments []*core.VolumeAttachment
		for _, attachment := range attachments {
			if attachment.VolumeID == disk.SelfLink {
				diskAttachments = append(diskAttachments, &core.VolumeAttachment{
					InstanceID: attachment.InstanceID,
					DeviceName: attachment.DeviceName,
					Status:     attachment.Status,
					VolumeID:   strconv.FormatUint(disk.Id, 10),
				})
			}
		}
		volumeSD := &core.Volume{
			Name:             disk.Name,
			VolumeID:         strconv.FormatUint(disk.Id, 10),
			AvailabilityZone: disk.Zone,
			Status:           disk.Status,
			VolumeType:       disk.Kind,
			NetworkName:      disk.SelfLink,
			IOPS:             0,
			Size:             strconv.FormatInt(disk.SizeGb, 10),
			Attachments:      diskAttachments,
		}
		volumesSD = append(volumesSD, volumeSD)

	}
	return volumesSD, nil
}

func (d *driver) GetVolumeAttach(
	volumeID, instanceID string) ([]*core.VolumeAttachment, error) {
	log.WithField("provider", providerName).Debug("GetVolumeAttach")
	return nil, nil
}

func (d *driver) waitSnapshotComplete(snapshotID string) error {
	return nil
}

func (d *driver) waitVolumeComplete(volumeID string) error {
	return nil
}

func (d *driver) waitVolumeAttach(volumeID, instanceID string) error {
	return nil
}

func (d *driver) waitVolumeDetach(volumeID string) error {
	return nil
}

func (d *driver) RemoveVolume(volumeID string) error {
	return nil
}

func (d *driver) AttachVolume(
	runAsync bool,
	volumeID, instanceID string) ([]*core.VolumeAttachment, error) {
	log.WithField("provider", providerName).Debug("AttachVolume")
	return nil, nil

}

func (d *driver) DetachVolume(
	runAsync bool,
	volumeID, blank string) error {
	log.WithField("provider", providerName).Debug("DetachVolume")
	return nil
}

func (d *driver) CopySnapshot(runAsync bool,
	volumeID, snapshotID, snapshotName, destinationSnapshotName,
	destinationRegion string) (*core.Snapshot, error) {
	log.WithField("provider", providerName).Debug("CopySnapshot")
	return nil, nil
}

func configRegistration() *config.Registration {
	r := config.NewRegistration("Google GCE")
	r.Key(config.String, "", "", "", "gce.zone")
	r.Key(config.String, "", "", "", "gce.project")
	r.Key(config.String, "", "", "", "gce.keyfile")
	return r
}
<commit_msg># This is a combination of 2 commits. # The first commit's message is:<commit_after>package gce

import (
	"bytes"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/emccode/rexray/core"
	"github.com/emccode/rexray/core/config"
	"github.com/emccode/rexray/core/errors"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/compute/v1"
	"io/ioutil"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const providerName = "gce"
func ef() errors.Fields {
	return errors.Fields{
		"provider": providerName,
	}
}

func eff(fields errors.Fields) map[string]interface{} {
	errFields := map[string]interface{}{
		"provider": providerName,
	}
	if fields != nil {
		for k, v := range fields {
			errFields[k] = v
		}
	}
	return errFields
}

// The GCE storage driver.
type driver struct {
	currentInstanceId string
	client            *compute.Service
	r                 *core.RexRay
	zone              string
	project           string
}

func init() {
	core.RegisterDriver(providerName, newDriver)
	config.Register(configRegistration())
}

func newDriver() core.Driver {
	return &driver{}
}

func (d *driver) Init(r *core.RexRay) error {
	d.r = r

	var err error

	d.zone = d.r.Config.GetString("gce.zone")
	d.project = d.r.Config.GetString("gce.project")
	serviceAccountJSON, err := ioutil.ReadFile(d.r.Config.GetString("gce.keyfile"))
	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not read service account credentials file, %s => {%s}", d.r.Config.GetString("gce.keyfile"), err)
		return err
	}

	config, err := google.JWTConfigFromJSON(serviceAccountJSON,
		compute.ComputeScope,
	)
	client, err := compute.New(config.Client(context.Background()))

	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not create compute client => {%s}", err)
		return err
	}
	d.client = client
	instanceId, err := getCurrentInstanceId()
	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not get current  instance => {%s}", err)
		return err
	}
	d.currentInstanceId = instanceId
	log.WithField("provider", providerName).Info("storage driver initialized")
	return nil
}

func getCurrentInstanceId() (string, error) {
	conn, err := net.DialTimeout("tcp", "metadata.google.internal:80", 50 * time.Millisecond)
	if err != nil {
		return "", err
	}
	defer conn.Close()

	url := "http://metadata.google.internal/computeMetadata/v1/instance/id"
	client := &http.Client{}
	req, _ := http.NewRequest("GET", url, nil)
	req.Header.Set("Metadata-Flavor", "Google")
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("Error: %v\n", err)
	}

	defer resp.Body.Close()

	buf := new(bytes.Buffer)
	buf.ReadFrom(resp.Body)
	s := buf.String()
	return s, nil
}

func (d *driver) Name() string {
	return providerName
}

func (d *driver) GetVolumeMapping() ([]*core.BlockDevice, error) {
	log.WithField("provider", providerName).Debug("GetVolumeMapping")

	diskMap := make(map[string]*compute.Disk)
	disks, err := d.client.Disks.List(d.project, d.zone).Do()
	if err != nil {
		return []*core.BlockDevice{}, err
	}
	for _, disk := range disks.Items {
		diskMap[disk.SelfLink] = disk
	}

	instances, err := d.client.Instances.List(d.project, d.zone).Do()
	if err != nil {
		return []*core.BlockDevice{}, err
	}
	var ret []*core.BlockDevice
	for _, instance := range instances.Items {
		for _, disk := range instance.Disks {
			ret = append(ret, &core.BlockDevice{
				ProviderName: "gce",
				InstanceID:   instance.Name,
				VolumeID:     diskMap[disk.Source].Name,
				DeviceName:   fmt.Sprintf("/dev/disk/by-id/google-%s", disk.DeviceName),
				Region:       diskMap[disk.Source].Zone,
				Status:       diskMap[disk.Source].Status,
				NetworkName:  disk.Source,
			})

		}
	}
	return ret, nil
}

func (d *driver) GetInstance() (*core.Instance, error) {
	log.WithField("provider", providerName).Debug("GetInstance")
	query := d.client.Instances.List(d.project, d.zone)
	instances, err := query.Do()
	if err != nil {
		return nil, err
	}
	for _, instance := range instances.Items {
		if strconv.FormatUint(instance.Id, 10) == d.currentInstanceId {
			return &core.Instance{
				ProviderName: "gce",
				InstanceID:   instance.Name,
				Region:       instance.Zone,
				Name:         instance.Name,
			}, nil
		}

	}
	return nil, nil
}

func (d *driver) CreateSnapshot(
runAsync bool,
snapshotName, volumeID, description string) ([]*core.Snapshot, error) {

	log.WithField("provider", providerName).Debug("CreateSnapshot")
	return nil, nil

}

func (d *driver) GetSnapshot(
volumeID, snapshotID, snapshotName string) ([]*core.Snapshot, error) {

	log.WithField("provider", providerName).Debug("GetSnapshot")
	return nil, nil
}

func (d *driver) RemoveSnapshot(snapshotID string) error {
	log.WithField("provider", providerName).Debug("RemoveSnapshot")
	return nil
}
func getLocalDevices() (deviceNames []string, err error) {
	file := "/proc/partitions"
	contentBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return []string{}, err
	}

	content := string(contentBytes)

	lines := strings.Split(content, "\n")
	for _, line := range lines[2:] {
		fields := strings.Fields(line)
		if len(fields) == 4 {
			deviceNames = append(deviceNames, fields[3])
		}
	}

	return deviceNames, nil
}

func (d *driver) GetDeviceNextAvailable() (string, error) {
	letters := []string{
		"a", "b", "c", "d", "e", "f", "g", "h",
		"i", "j", "k", "l", "m", "n", "o", "p"}

	blockDeviceNames := make(map[string]bool)

	blockDeviceMapping, err := d.GetVolumeMapping()
	if err != nil {
		return "", err
	}

	for _, blockDevice := range blockDeviceMapping {
		re, _ := regexp.Compile(`^/dev/sd([a-z])`)
		res := re.FindStringSubmatch(blockDevice.DeviceName)
		if len(res) > 0 {
			blockDeviceNames[res[1]] = true
		}
	}

	localDevices, err := getLocalDevices()
	if err != nil {
		return "", err
	}

	for _, localDevice := range localDevices {
		re, _ := regexp.Compile(`^sd([a-z])`)
		res := re.FindStringSubmatch(localDevice)
		if len(res) > 0 {
			blockDeviceNames[res[1]] = true
		}
	}

	for _, letter := range letters {
		if !blockDeviceNames[letter] {
			nextDeviceName := "/dev/sd" + letter
			log.Println("Got next device name: " + nextDeviceName)
			return nextDeviceName, nil
		}
	}
	return "", errors.New("No available device")
}
func (d *driver) waitUntilOperationIsFinished(operation *compute.Operation) error {
	opName := operation.Name
	OpLoop:
	for {
		time.Sleep(100 * time.Millisecond)
		op, err := d.client.ZoneOperations.Get(d.project, d.zone, opName).Do()
		if err != nil {
			return err
		}
		switch op.Status {
		case "PENDING", "RUNNING":
			continue
		case "DONE":
			if op.Error != nil {
				return err
			}
			break OpLoop
		default:
			log.WithField("provider", providerName).Fatalf("Unknown status %q: %+v", op.Status, op)
			return nil
		}
	}
	return nil
}
func (d *driver) CreateVolume(
runAsync bool, volumeName, volumeID, snapshotID, volumeType string,
IOPS, size int64, availabilityZone string) (*core.Volume, error) {
	log.WithFields(log.Fields{
		"provider":         providerName,
		"volumeName":       volumeName,
		"volumeID":         volumeID,
		"snapshotID":       snapshotID,
		"volumeType":       volumeType,
		"availabilityZone": availabilityZone}).Debug("CreateVolume")

	if availabilityZone == "" {
		availabilityZone = d.zone
	}
	disk := &compute.Disk{
		Name:   volumeName,
		Zone:   availabilityZone,
		Type:   "https://www.googleapis.com/compute/v1/projects/gce-dev-1060/zones/europe-west1-b/diskTypes/pd-standard",
		SizeGb: size,
	}
	createdVolume, err := d.client.Disks.Insert(d.project, d.zone, disk).Do()
	if err != nil {
		return nil, err
	}
	if !runAsync {
		err := d.waitUntilOperationIsFinished(createdVolume)
		if err != nil {
			return nil, err
		}

	}
	volume, err := d.GetVolume(volumeName, "")
	if err != nil {
		return nil, err
	}
	return volume[0], nil
}

func (d *driver) getVolumesAttachedToInstance(instances []*compute.Instance,volumeID string) []*core.VolumeAttachment {
	var attachments []*core.VolumeAttachment
	for _, instance := range instances {
		for _, disk := range instance.Disks {
			shortVolume := strings.Replace(disk.Source,
				fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/zones/%s/disks/", d.project, d.zone),
				"",
				-1)

			if volumeID !="" && shortVolume!=volumeID {
				continue
			}
			attachments = append(attachments, d.convertGCEAttachedDisk(instance.Name, disk))
		}
	}

	return attachments
}
func (d *driver) GetVolume(
volumeID, volumeName string) ([]*core.Volume, error) {
	log.WithField("provider", providerName).Debugf("GetVolume :%s %s", volumeID, volumeName)

	instances, err := d.client.Instances.List(d.project, d.zone).Do()
	if err != nil {
		return []*core.Volume{}, err
	}
	attachments := d.getVolumesAttachedToInstance(instances.Items,"")

	query := d.client.Disks.List(d.project, d.zone)
	if volumeID != "" {
		query.Filter(fmt.Sprintf("name eq %s", volumeID))
	}
	if volumeName != "" {
		query.Filter(fmt.Sprintf("name eq %s", volumeName))
	}
	disks, err := query.Do()
	if err != nil {
		return []*core.Volume{}, err
	}
	var volumesSD []*core.Volume
	for _, disk := range disks.Items {
		var diskAttachments []*core.VolumeAttachment
		for _, attachment := range attachments {
			if attachment.VolumeID == disk.SelfLink {
				diskAttachments = append(diskAttachments, &core.VolumeAttachment{
					InstanceID: attachment.InstanceID,
					DeviceName: fmt.Sprintf("/dev/disk/by-id/google-%s", attachment.DeviceName),
					Status:     attachment.Status,
					VolumeID:   strconv.FormatUint(disk.Id, 10),
				})
			}
		}
		volumeSD := &core.Volume{
			Name:             disk.Name,
			VolumeID:         disk.Name,
			AvailabilityZone: disk.Zone,
			Status:           disk.Status,
			VolumeType:       disk.Kind,
			NetworkName:      disk.SelfLink,
			IOPS:             0,
			Size:             strconv.FormatInt(disk.SizeGb, 10),
			Attachments:      diskAttachments,
		}
		volumesSD = append(volumesSD, volumeSD)

	}
	return volumesSD, nil
}

func (d *driver) convertGCEAttachedDisk(instanceName string, disk *compute.AttachedDisk) *core.VolumeAttachment {
	return &core.VolumeAttachment{
		InstanceID: instanceName,
		DeviceName: fmt.Sprintf("/dev/disk/by-id/google-%s", disk.DeviceName),
		Status:     disk.Mode,
		VolumeID:   disk.Source,
	}
}

func (d *driver) GetVolumeAttach(
volumeID, instanceID string) ([]*core.VolumeAttachment, error) {
	log.WithField("provider", providerName).Debugf("GetVolumeAttach :%s %s", volumeID, instanceID)
	query := d.client.Instances.List(d.project, d.zone)
	if instanceID != "" {
		query.Filter(fmt.Sprintf("name eq %s", instanceID))
	}
	instances, err := query.Do()
	if err != nil {
		return []*core.VolumeAttachment{}, err
	}
	return d.getVolumesAttachedToInstance(instances.Items,volumeID),nil;
}

func (d *driver) RemoveVolume(volumeID string) error {
	log.WithField("provider", providerName).Debugf("RemoveVolume :%s", volumeID)
	_, err := d.client.Disks.Delete(d.project, d.zone, volumeID).Do()
	return err

}

func (d *driver) AttachVolume(
runAsync bool,
volumeID, instanceID string) ([]*core.VolumeAttachment, error) {
	if instanceID == "" {
		instanceID = d.currentInstanceId
	}
	instance, err := d.GetInstance()
	if err != nil {
		return nil, err
	}
	instanceID = instance.Name
	log.WithField("provider", providerName).Debugf("AttachVolume %s %s", volumeID, instance.Name)
	query := d.client.Disks.List(d.project, d.zone)
	query.Filter(fmt.Sprintf("name eq %s", volumeID))
	disks, err := query.Do()
	if err != nil {
		return nil, err
	}
	if len(disks.Items) != 1 {
		return nil, errors.New("No available device")
	}

	disk := &compute.AttachedDisk{
		AutoDelete: false,
		Boot:       false,
		Source:     disks.Items[0].SelfLink,
	}
	operation, err := d.client.Instances.AttachDisk(d.project, d.zone, instanceID, disk).Do()
	if err != nil {
		return nil, err
	}
	if !runAsync {
		err := d.waitUntilOperationIsFinished(operation)
		if err != nil {
			return nil, err
		}
	}

	return d.GetVolumeAttach(volumeID, instanceID)

}

func (d *driver) DetachVolume(
runAsync bool,
volumeID, blank string) error {
	instance, err := d.GetInstance()
	if err != nil {
		return err
	}
	instanceID := instance.Name
	log.WithField("provider", providerName).Debugf("DetachVolume %s %s", volumeID, instance.Name)
	query := d.client.Disks.List(d.project, d.zone)
	query.Filter(fmt.Sprintf("name eq %s", volumeID))
	disks, err := query.Do()
	if err != nil {
		return err
	}
	if len(disks.Items) != 1 {
		return errors.New("No available device")
	}
	attachements, err := d.GetVolumeAttach(volumeID, instanceID)
	for _, attachement := range attachements {
		targetVolumeId := strings.Replace(attachement.VolumeID,
			fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/zones/%s/disks/", d.project, d.zone),
			"",
			-1)
		if (targetVolumeId == volumeID) {
			diskName := strings.Replace(attachement.DeviceName, "/dev/disk/by-id/google-", "", -1)
			operation, err := d.client.Instances.DetachDisk(d.project, d.zone, instanceID, diskName).Do()
			if err != nil {
				return err
			}
			if !runAsync {
				err := d.waitUntilOperationIsFinished(operation)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func (d *driver) CopySnapshot(runAsync bool,
volumeID, snapshotID, snapshotName, destinationSnapshotName,
destinationRegion string) (*core.Snapshot, error) {
	log.WithField("provider", providerName).Debug("CopySnapshot")
	return nil, nil
}

func configRegistration() *config.Registration {
	r := config.NewRegistration("Google GCE")
	r.Key(config.String, "", "", "", "gce.zone")
	r.Key(config.String, "", "", "", "gce.project")
	r.Key(config.String, "", "", "", "gce.keyfile")
	return r
}
<|endoftext|>


<commit_before>package x509

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"github.com/mandelsoft/spiff/yaml"
	"math/big"
	"net"
	"time"

	. "github.com/mandelsoft/spiff/dynaml"
)

const F_Cert = "x509cert"

func init() {
	RegisterFunction(F_Cert, func_x509cert)
}

//  one map argument with fields
//   usage:        []string
//   organization: []string
//   country: 	   []string (optional)
//   commonName:   string   (optional)
//   validFrom :   string/date   (optional)
//   validity :    int    (hours, optional)
//   isCA:         boolean  (optional)
//   hosts:        []string (optional)
//   privateKey:   string
//   publicKey:    string
//
//   caCert:       string   (optional)
//   caPrivateKey: string   (optional)
//

func func_x509cert(arguments []interface{}, binding Binding) (interface{}, EvaluationInfo, bool) {
	var err error
	info := DefaultInfo()

	if len(arguments) != 1 {
		return info.Error("invalid argument count for %s(<map>)", F_Cert)
	}
	fields, ok := arguments[0].(map[string]yaml.Node)
	if !ok {
		return info.Error("argument for %s must be a map (found %T)", F_Cert, arguments[0])
	}

	isCA, err := getDefaultedBoolField(fields, "isCA", false)
	if err != nil {
		return info.Error(err)
	}

	orgs, err := getStringListField(fields, "organization")
	if err != nil {
		return info.Error(err)
	}

	cn, err := getDefaultedStringField(fields, "commonName", "")
	if err != nil {
		return info.Error(err)
	}

	countries, err := getDefaultedStringListField(fields, "country", nil)
	if err != nil {
		return info.Error(err)
	}

	usages, err := getStringListField(fields, "usage")
	if err != nil {
		return info.Error(err)
	}

	validity, err := getDefaultedIntField(fields, "validity", 24*365)
	if err != nil {
		return info.Error(err)
	}

	hosts, err := getDefaultedStringListField(fields, "hosts", nil)
	if err != nil {
		return info.Error(err)
	}

	privKey, err := getDefaultedStringField(fields, "privateKey", "")
	if err != nil {
		return info.Error(err)
	}
	var priv interface{}
	if privKey != "" {
		priv, err = ParsePrivateKey(privKey)
		if err != nil {
			return info.Error(err)
		}
	}

	pubKey, err := getDefaultedStringField(fields, "publicKey", "")
	if err != nil {
		return info.Error(err)
	}
	var pub interface{}
	if pubKey != "" {
		pub, err = ParsePublicKey(pubKey)
		if err != nil {
			return info.Error(err)
		}
	}

	if pub == nil {
		if priv == nil {
			return info.Error("one of 'publicKey' or 'privKey' must be given")
		}
		pub = publicKey(priv)
	}

	caCert, err := getDefaultedStringField(fields, "caCert", "")
	if err != nil {
		return info.Error(err)
	}
	var ca *x509.Certificate
	if caCert != "" {
		ca, err = ParseCertificate(caCert)
		if err != nil {
			return info.Error("invalid ca certificate: %s", err)
		}
	}

	caPrivateKey, err := getDefaultedStringField(fields, "caPrivateKey", "")
	if err != nil {
		return info.Error(err)
	}
	var capriv = priv
	if caPrivateKey != "" {
		if ca != nil {
			capriv, err = ParsePrivateKey(caPrivateKey)
			if err != nil {
				return info.Error(err)
			}
		}
	} else {
		if ca != nil {
			return info.Error("private key for ca required")
		}
	}
	if capriv == nil {
		return info.Error("private key for self-signed certificate required")
	}

	var notBefore time.Time
	validFrom, err := getDefaultedStringField(fields, "validFrom", "")
	if err != nil {
		return info.Error(err)
	}
	if validFrom == "" {
		notBefore = time.Now()
	} else {
		notBefore, err = time.Parse("Jan 2 15:04:05 2006", validFrom)
		if err != nil {
			return info.Error("invalid validFrom fields: %s", err)
		}
	}

	notAfter := notBefore.Add(time.Duration(validity) * time.Hour)

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return info.Error("failed to generate serial number: %s", err)
	}

	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: orgs,
			CommonName:   cn,
			Country:      countries,
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	if ca == nil {
		ca = template
	}

	for _, u := range usages {
		k := ParseKeyUsage(u)
		if k == nil {
			return info.Error("invalid usage key %q", u)
		}
		k.AddTo(template)
	}

	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, h)
		}
	}

	if isCA || (template.KeyUsage&x509.KeyUsageCertSign) != 0 {
		template.IsCA = true
		template.KeyUsage |= x509.KeyUsageCertSign
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, template, ca, pub, capriv)
	if err != nil {
		return info.Error("Failed to create certificate: %s", err)
	}

	var b bytes.Buffer
	writer := bufio.NewWriter(&b)

	if err := pem.Encode(writer, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes}); err != nil {
		return info.Error("failed to write certificate pem block: %s", err)
	}
	writer.Flush()
	return b.String(), info, true
}
<commit_msg>fix x509<commit_after>package x509

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"github.com/mandelsoft/spiff/yaml"
	"math/big"
	"net"
	"time"

	. "github.com/mandelsoft/spiff/dynaml"
)

const F_Cert = "x509cert"

func init() {
	RegisterFunction(F_Cert, func_x509cert)
}

//  one map argument with fields
//   usage:        []string
//   organization: []string
//   country: 	   []string (optional)
//   commonName:   string   (optional)
//   validFrom :   string/date   (optional)
//   validity :    int    (hours, optional)
//   isCA:         boolean  (optional)
//   hosts:        []string (optional)
//   privateKey:   string
//   publicKey:    string
//
//   caCert:       string   (optional)
//   caPrivateKey: string   (optional)
//

func func_x509cert(arguments []interface{}, binding Binding) (interface{}, EvaluationInfo, bool) {
	var err error
	info := DefaultInfo()

	if len(arguments) != 1 {
		return info.Error("invalid argument count for %s(<map>)", F_Cert)
	}
	fields, ok := arguments[0].(map[string]yaml.Node)
	if !ok {
		return info.Error("argument for %s must be a map (found %T)", F_Cert, arguments[0])
	}

	isCA, err := getDefaultedBoolField(fields, "isCA", false)
	if err != nil {
		return info.Error(err)
	}

	orgs, err := getStringListField(fields, "organization")
	if err != nil {
		return info.Error(err)
	}

	cn, err := getDefaultedStringField(fields, "commonName", "")
	if err != nil {
		return info.Error(err)
	}

	countries, err := getDefaultedStringListField(fields, "country", nil)
	if err != nil {
		return info.Error(err)
	}

	usages, err := getStringListField(fields, "usage")
	if err != nil {
		return info.Error(err)
	}

	validity, err := getDefaultedIntField(fields, "validity", 24*365)
	if err != nil {
		return info.Error(err)
	}

	hosts, err := getDefaultedStringListField(fields, "hosts", nil)
	if err != nil {
		return info.Error(err)
	}

	privKey, err := getDefaultedStringField(fields, "privateKey", "")
	if err != nil {
		return info.Error(err)
	}
	var priv interface{}
	if privKey != "" {
		priv, err = ParsePrivateKey(privKey)
		if err != nil {
			return info.Error(err)
		}
	}

	pubKey, err := getDefaultedStringField(fields, "publicKey", "")
	if err != nil {
		return info.Error(err)
	}
	var pub interface{}
	if pubKey != "" {
		pub, err = ParsePublicKey(pubKey)
		if err != nil {
			return info.Error(err)
		}
	}

	if pub == nil {
		if priv == nil {
			return info.Error("one of 'publicKey' or 'privKey' must be given")
		}
		pub = publicKey(priv)
	}

	caCert, err := getDefaultedStringField(fields, "caCert", "")
	if err != nil {
		return info.Error(err)
	}
	var ca *x509.Certificate
	if caCert != "" {
		ca, err = ParseCertificate(caCert)
		if err != nil {
			return info.Error("invalid ca certificate: %s", err)
		}
	}

	caPrivateKey, err := getDefaultedStringField(fields, "caPrivateKey", "")
	if err != nil {
		return info.Error(err)
	}
	var capriv = priv
	if caPrivateKey != "" {
		if ca != nil {
			capriv, err = ParsePrivateKey(caPrivateKey)
			if err != nil {
				return info.Error(err)
			}
		}
	} else {
		if ca != nil {
			return info.Error("private key for ca required")
		}
	}
	if capriv == nil {
		return info.Error("private key for self-signed certificate required")
	}

	var notBefore time.Time
	validFrom, err := getDefaultedStringField(fields, "validFrom", "")
	if err != nil {
		return info.Error(err)
	}
	if validFrom == "" {
		notBefore = time.Now()
	} else {
		notBefore, err = time.Parse("Jan 2 15:04:05 2006", validFrom)
		if err != nil {
			return info.Error("invalid validFrom fields: %s", err)
		}
	}

	notAfter := notBefore.Add(time.Duration(validity) * time.Hour)

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return info.Error("failed to generate serial number: %s", err)
	}

	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: orgs,
			CommonName:   cn,
			Country:      countries,
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              0,
		ExtKeyUsage:           []x509.ExtKeyUsage{},
		BasicConstraintsValid: true,
	}

	if ca == nil {
		ca = template
	}

	for _, u := range usages {
		k := ParseKeyUsage(u)
		if k == nil {
			return info.Error("invalid usage key %q", u)
		}
		k.AddTo(template)
	}

	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, h)
		}
	}

	if isCA || (template.KeyUsage&x509.KeyUsageCertSign) != 0 {
		template.IsCA = true
		template.KeyUsage |= x509.KeyUsageCertSign
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, template, ca, pub, capriv)
	if err != nil {
		return info.Error("Failed to create certificate: %s", err)
	}

	var b bytes.Buffer
	writer := bufio.NewWriter(&b)

	if err := pem.Encode(writer, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes}); err != nil {
		return info.Error("failed to write certificate pem block: %s", err)
	}
	writer.Flush()
	return b.String(), info, true
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math/rand"
	"time"
)

func pirate(name string) {
	cnt := 0
	for {
		cnt++
		time.Sleep((200 + time.Duration(rand.Intn(1300))) * time.Millisecond)
		fmt.Printf("%v says: heyo for %vth time! \n", name, cnt)
	}
}

func main() {
	for i := 65; i < 90; i++ {
		go pirate(fmt.Sprintf("Pirate%v", string(rune(i))))
	}
	time.Sleep(5000 * time.Millisecond)
}
<commit_msg>Showcase some channel communication<commit_after>package main

import (
	"fmt"
	"math/rand"
	"time"
)

func pirate(name string, out chan string) {
	cnt := 0
	for {
		cnt++
		time.Sleep((time.Duration(rand.Intn(1300))) * time.Millisecond)
		if cnt > 3 {
			out <- fmt.Sprintf("%v is tired of this shit and leaves\n", name)
			return
		}
		res := fmt.Sprintf("%v says: heyo for %vth time! \n", name, cnt)
		out <- res
	}
}

func main() {
	log := make(chan string)
	for i := 65; i < 68; i++ {
		go pirate(fmt.Sprintf("Pirate %v", string(rune(i))), log)
	}

	go func() {
		for l := range log {
			fmt.Printf(l)
		}
	}()
	time.Sleep(10000 * time.Millisecond)
}
<|endoftext|>


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

package lxd_client

import (
	"crypto/x509"
	"io"
	"os"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"
)

// These interfaces facilitate mocking out the LXD API during tests.
// See https://github.com/lxc/lxd/blob/master/client.go
// and https://github.com/lxc/lxd/blob/master/specs/rest-api.md.

type rawClientWrapper interface {
	rawServerMethods
	rawImageMethods
	rawAliasMethods
	rawContainerMethods
	rawProfileMethods
}

type rawServerMethods interface {
	Finger() error
	ServerStatus() (*shared.ServerState, error)
	GetServerConfigString() ([]string, error)
	SetServerConfig(key string, value string) (*lxd.Response, error)

	// connection
	WaitFor(waitURL string) (*shared.Operation, error)
	WaitForSuccess(waitURL string) error

	// auth
	UserAuthServerCert(name string, acceptCert bool) error
	CertificateList() ([]shared.CertInfo, error)
	AddMyCertToServer(pwd string) error
	CertificateAdd(cert *x509.Certificate, name string) error
	CertificateRemove(fingerprint string) error
	AmTrusted() bool
}

type rawImageMethods interface {
	ListImages() ([]shared.ImageInfo, error)
	GetImageInfo(image string) (*shared.ImageInfo, error)
	PutImageProperties(name string, p shared.ImageProperties) error

	CopyImage(image string, dest *lxd.Client, copy_aliases bool, aliases []string, public bool) error
	PostImage(imageFile string, rootfsFile string, properties []string, public bool, aliases []string) (string, error)
	ExportImage(image string, target string) (*lxd.Response, string, error)
	DeleteImage(image string) error

	ImageFromContainer(cname string, public bool, aliases []string, properties map[string]string) (string, error)
}

type rawAliasMethods interface {
	ListAliases() ([]shared.ImageAlias, error)
	IsAlias(alias string) (bool, error)

	GetAlias(alias string) string
	PostAlias(alias string, desc string, target string) error
	DeleteAlias(alias string) error
}

type rawContainerMethods interface {
	ListContainers() ([]shared.ContainerInfo, error)
	Rename(name string, newName string) (*lxd.Response, error)
	ContainerStatus(name string, showLog bool) (*shared.ContainerState, error)

	Init(name string, imgremote string, image string, profiles *[]string, ephem bool) (*lxd.Response, error)
	LocalCopy(source string, name string, config map[string]string, profiles []string) (*lxd.Response, error)
	MigrateTo(container string) (*lxd.Response, error)
	MigrateFrom(name string, operation string, secrets map[string]string, config map[string]string, profiles []string, baseImage string) (*lxd.Response, error)
	Action(name string, action shared.ContainerAction, timeout int, force bool) (*lxd.Response, error)
	Delete(name string) (*lxd.Response, error)

	// exec
	Exec(name string, cmd []string, env map[string]string, stdin *os.File, stdout *os.File, stderr *os.File) (int, error)

	// files
	PushFile(container string, p string, gid int, uid int, mode os.FileMode, buf io.ReadSeeker) error
	PullFile(container string, p string) (int, int, os.FileMode, io.ReadCloser, error)

	// config
	GetContainerConfig(container string) ([]string, error)
	SetContainerConfig(container, key, value string) (*lxd.Response, error)
	UpdateContainerConfig(container string, st shared.BriefContainerState) error

	// devices
	ContainerListDevices(container string) ([]string, error)
	ContainerDeviceDelete(container, devname string) (*lxd.Response, error)
	ContainerDeviceAdd(container, devname, devtype string, props []string) (*lxd.Response, error)

	// snapshots
	RestoreSnapshot(container string, snapshotName string, stateful bool) (*lxd.Response, error)
	Snapshot(container string, snapshotName string, stateful bool) (*lxd.Response, error)
	ListSnapshots(container string) ([]string, error)
}

type rawProfileMethods interface {
	ListProfiles() ([]string, error)

	ProfileCreate(p string) error
	PutProfile(name string, profile shared.ProfileConfig) error
	ProfileCopy(name, newname string, dest *lxd.Client) error
	ProfileDelete(p string) error

	// apply
	ApplyProfile(container, profile string) (*lxd.Response, error)

	// config
	ProfileConfig(name string) (*shared.ProfileConfig, error)
	GetProfileConfig(profile string) (map[string]string, error)
	SetProfileConfigItem(profile, key, value string) error

	// devices
	ProfileListDevices(profile string) ([]string, error)
	ProfileDeviceDelete(profile, devname string) (*lxd.Response, error)
	ProfileDeviceAdd(profile, devname, devtype string, props []string) (*lxd.Response, error)
}
<commit_msg>Clean up the rawClientWrapper interface.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package lxd_client

import (
	"crypto/x509"
	"io"
	"os"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"
)

// These interfaces facilitate mocking out the LXD API during tests.
// See https://github.com/lxc/lxd/blob/master/client.go
// and https://github.com/lxc/lxd/blob/master/specs/rest-api.md.

// TODO(ericsnow) Move these to a test suite.
var (
	_ rawClientWrapper     = (*lxd.Client)(nil)
	_ rawClientWrapperFull = (*lxd.Client)(nil)
)

// rawClientWrapper exposes the methods of lxd.Client required for this
// high-level client package.
type rawClientWrapper interface {
	// server methods
	WaitForSuccess(waitURL string) error

	// container methods
	ListContainers() ([]shared.ContainerInfo, error)
	ContainerStatus(name string, showLog bool) (*shared.ContainerState, error)
	Init(name string, imgremote string, image string, profiles *[]string, ephem bool) (*lxd.Response, error)
	Action(name string, action shared.ContainerAction, timeout int, force bool) (*lxd.Response, error)
	Delete(name string) (*lxd.Response, error)
}

// rawClientWrapperFull exposes all methods of lxd.Client.
type rawClientWrapperFull interface {
	rawServerMethods
	rawImageMethods
	rawAliasMethods
	rawContainerMethods
	rawProfileMethods
}

type rawServerMethods interface {
	// info
	Finger() error
	ServerStatus() (*shared.ServerState, error)

	// config
	GetServerConfigString() ([]string, error)
	SetServerConfig(key string, value string) (*lxd.Response, error)

	// connection
	WaitFor(waitURL string) (*shared.Operation, error)
	WaitForSuccess(waitURL string) error

	// auth
	UserAuthServerCert(name string, acceptCert bool) error
	CertificateList() ([]shared.CertInfo, error)
	AddMyCertToServer(pwd string) error
	CertificateAdd(cert *x509.Certificate, name string) error
	CertificateRemove(fingerprint string) error
	AmTrusted() bool
}

type rawImageMethods interface {
	// info/meta
	ListImages() ([]shared.ImageInfo, error)
	GetImageInfo(image string) (*shared.ImageInfo, error)
	PutImageProperties(name string, p shared.ImageProperties) error

	// image data (create, upload, download, destroy)
	CopyImage(image string, dest *lxd.Client, copy_aliases bool, aliases []string, public bool) error
	ImageFromContainer(cname string, public bool, aliases []string, properties map[string]string) (string, error)
	PostImage(imageFile string, rootfsFile string, properties []string, public bool, aliases []string) (string, error)
	ExportImage(image string, target string) (*lxd.Response, string, error)
	DeleteImage(image string) error
}

type rawAliasMethods interface {
	// info
	ListAliases() ([]shared.ImageAlias, error)
	IsAlias(alias string) (bool, error)

	// alias data (upload, download, destroy)
	PostAlias(alias string, desc string, target string) error
	GetAlias(alias string) string
	DeleteAlias(alias string) error
}

type rawContainerMethods interface {
	// info/meta
	ListContainers() ([]shared.ContainerInfo, error)
	//Rename(name string, newName string) (*lxd.Response, error)
	ContainerStatus(name string, showLog bool) (*shared.ContainerState, error)

	// container data (create, actions, destroy)
	Init(name string, imgremote string, image string, profiles *[]string, ephem bool) (*lxd.Response, error)
	LocalCopy(source string, name string, config map[string]string, profiles []string) (*lxd.Response, error)
	MigrateTo(container string) (*lxd.Response, error)
	MigrateFrom(name string, operation string, secrets map[string]string, config map[string]string, profiles []string, baseImage string) (*lxd.Response, error)
	Action(name string, action shared.ContainerAction, timeout int, force bool) (*lxd.Response, error)
	Delete(name string) (*lxd.Response, error)

	// exec
	Exec(name string, cmd []string, env map[string]string, stdin *os.File, stdout *os.File, stderr *os.File) (int, error)

	// files
	PushFile(container string, p string, gid int, uid int, mode os.FileMode, buf io.ReadSeeker) error
	PullFile(container string, p string) (int, int, os.FileMode, io.ReadCloser, error)

	// config
	GetContainerConfig(container string) ([]string, error)
	SetContainerConfig(container, key, value string) (*lxd.Response, error)
	UpdateContainerConfig(container string, st shared.BriefContainerState) error

	// devices
	ContainerListDevices(container string) ([]string, error)
	ContainerDeviceDelete(container, devname string) (*lxd.Response, error)
	ContainerDeviceAdd(container, devname, devtype string, props []string) (*lxd.Response, error)

	// snapshots
	RestoreSnapshot(container string, snapshotName string, stateful bool) (*lxd.Response, error)
	Snapshot(container string, snapshotName string, stateful bool) (*lxd.Response, error)
	ListSnapshots(container string) ([]string, error)
}

type rawProfileMethods interface {
	// info
	ListProfiles() ([]string, error)

	// profile data (create, upload, destroy)
	ProfileCreate(p string) error
	ProfileCopy(name, newname string, dest *lxd.Client) error
	PutProfile(name string, profile shared.ProfileConfig) error
	ProfileDelete(p string) error

	// apply
	ApplyProfile(container, profile string) (*lxd.Response, error)

	// config
	ProfileConfig(name string) (*shared.ProfileConfig, error)
	GetProfileConfig(profile string) (map[string]string, error)
	SetProfileConfigItem(profile, key, value string) error

	// devices
	ProfileListDevices(profile string) ([]string, error)
	ProfileDeviceDelete(profile, devname string) (*lxd.Response, error)
	ProfileDeviceAdd(profile, devname, devtype string, props []string) (*lxd.Response, error)
}
<|endoftext|>


<commit_before>package main

import (
	"reflect"
	"testing"
)

func TestMove(t *testing.T) {
	s := []rune("aaaaa bbbbb ccccc")
	tests := []struct {
		input int
		want  int
	}{
		{0, 0},
		{1, 1},
		{10, 10},
		{len(s), len(s)},
		{-1, 0},
		{len(s) + 1, len(s)},
	}
	for _, test := range tests {
		e := &basicEditor{buf: s}
		e.move(test.input)
		if e.pos != test.want {
			t.Errorf("move(%v): got %v, want %v", test.input, e.pos, test.want)
		}
	}
}

func TestInsert(t *testing.T) {
	tests := []struct {
		initial basicEditor
		input   []rune
		at      int
		want    basicEditor
	}{
		{
			initial: basicEditor{buf: []rune(""), pos: 0},
			input:   []rune("aaa"),
			at:      0,
			want:    basicEditor{buf: []rune("aaa"), pos: 3},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 2},
			input:   []rune("aaa"),
			at:      -1,
			want:    basicEditor{buf: []rune("aaaAAA"), pos: 5},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 1},
			input:   []rune("aaa"),
			at:      2,
			want:    basicEditor{buf: []rune("AAaaaA"), pos: 1},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 3},
			input:   []rune("aaa"),
			at:      10,
			want:    basicEditor{buf: []rune("AAAaaa"), pos: 6},
		},
	}
	for _, test := range tests {
		test.initial.insert(test.input, test.at)
		if !reflect.DeepEqual(test.initial, test.want) {
			t.Errorf("insert(%v, %v): got %v, want %v", test.input, test.at, test.initial, test.want)
		}
	}
}

func TestDelete(t *testing.T) {
	tests := []struct {
		initial basicEditor
		from    int
		to      int
		want    basicEditor
	}{
		{
			initial: basicEditor{buf: []rune(""), pos: 0},
			from:    -1,
			to:      1,
			want:    basicEditor{buf: []rune(""), pos: 0},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 2},
			from:    1,
			to:      2,
			want:    basicEditor{buf: []rune("AA"), pos: 1},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 1},
			from:    2,
			to:      0,
			want:    basicEditor{buf: []rune("A"), pos: 0},
		},
		{
			initial: basicEditor{buf: []rune("AAAABBCCaaaa"), pos: 3},
			from:    4,
			to:      8,
			want:    basicEditor{buf: []rune("AAAAaaaa"), pos: 3},
		},
		{
			initial: basicEditor{buf: []rune(""), pos: 0},
			from:    -1,
			to:      -1,
			want:    basicEditor{buf: []rune(""), pos: 0},
		},
	}
	for _, test := range tests {
		test.initial.delete(test.from, test.to)
		if !reflect.DeepEqual(test.initial, test.want) {
			t.Errorf("delete(%v, %v): got %v, want %v", test.from, test.to, test.initial, test.want)
		}
	}
}

func TestSlice(t *testing.T) {
	tests := []struct {
		initial []rune
		from    int
		to      int
		want    []rune
	}{
		{
			initial: []rune(""),
			from:    -1,
			to:      1,
			want:    []rune(""),
		},
		{
			initial: []rune("ABC"),
			from:    -1,
			to:      1,
			want:    []rune("A"),
		},
		{
			initial: []rune("aaa bbb ccc"),
			from:    2,
			to:      9,
			want:    []rune("a bbb c"),
		},
		{
			initial: []rune("aaa x bbb"),
			from:    10,
			to:      4,
			want:    []rune("x bbb"),
		},
	}
	for _, test := range tests {
		e := &basicEditor{buf: test.initial}
		got := e.slice(test.from, test.to)
		if string(got) != string(test.want) {
			t.Errorf("slice(%v, %v): got %v, want %v", test.from, test.to, got, test.want)
		}
	}
}

func TestIndex(t *testing.T) {
	tests := []struct {
		initial []rune
		ch      rune
		start   int
		want    int
	}{
		{
			initial: []rune(""),
			ch:      'a',
			start:   0,
			want:    -1,
		},
		{
			initial: []rune("aaa"),
			ch:      'a',
			start:   1,
			want:    1,
		},
		{
			initial: []rune("abcde"),
			ch:      'e',
			start:   3,
			want:    4,
		},
		{
			initial: []rune("AA AA"),
			ch:      'A',
			start:   2,
			want:    3,
		},
	}
	for _, test := range tests {
		e := &basicEditor{buf: test.initial}
		got := e.index(test.ch, test.start)
		if got != test.want {
			t.Errorf("index(%v, %v): got %v, want %v", test.ch, test.start, got, test.want)
		}
	}
}

func TestLastIndex(t *testing.T) {
	tests := []struct {
		initial []rune
		ch      rune
		start   int
		want    int
	}{
		{
			initial: []rune(""),
			ch:      'a',
			start:   0,
			want:    -1,
		},
		{
			initial: []rune("aaa"),
			ch:      'a',
			start:   1,
			want:    0,
		},
		{
			initial: []rune("abcde"),
			ch:      'e',
			start:   3,
			want:    -1,
		},
		{
			initial: []rune("AA AA"),
			ch:      'A',
			start:   2,
			want:    1,
		},
	}
	for _, test := range tests {
		e := &basicEditor{buf: test.initial}
		got := e.lastIndex(test.ch, test.start)
		if got != test.want {
			t.Errorf("lastIndex(%v, %v): got %v, want %v", string(test.ch), test.start, got, test.want)
		}
	}
}

func TestReplace(t *testing.T) {
	tests := []struct {
		initial basicEditor
		input   []rune
		at      int
		want    basicEditor
	}{
		{
			initial: basicEditor{buf: []rune(""), pos: 0},
			input:   []rune("aaa"),
			at:      0,
			want:    basicEditor{buf: []rune("aaa"), pos: 0},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 2},
			input:   []rune("aaa"),
			at:      -1,
			want:    basicEditor{buf: []rune("aaA"), pos: 2},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 1},
			input:   []rune("aaa"),
			at:      2,
			want:    basicEditor{buf: []rune("AAaaa"), pos: 1},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 3},
			input:   []rune("aaa"),
			at:      10,
			want:    basicEditor{buf: []rune("AAA       aaa"), pos: 3},
		},
	}
	for _, test := range tests {
		test.initial.replace(test.input, test.at)
		if !reflect.DeepEqual(test.initial, test.want) {
			t.Errorf("replace(%v, %v): got %v, want %v", string(test.input), test.at, test.initial, test.want)
		}
	}
}
<commit_msg>Add a test case which makes test fail<commit_after>package main

import (
	"reflect"
	"testing"
)

func TestMove(t *testing.T) {
	s := []rune("aaaaa bbbbb ccccc")
	tests := []struct {
		input int
		want  int
	}{
		{0, 0},
		{1, 1},
		{10, 10},
		{len(s), len(s)},
		{-1, 0},
		{len(s) + 1, len(s)},
	}
	for _, test := range tests {
		e := &basicEditor{buf: s}
		e.move(test.input)
		if e.pos != test.want {
			t.Errorf("move(%v): got %v, want %v", test.input, e.pos, test.want)
		}
	}
}

func TestInsert(t *testing.T) {
	tests := []struct {
		initial basicEditor
		input   []rune
		at      int
		want    basicEditor
	}{
		{
			initial: basicEditor{buf: []rune(""), pos: 0},
			input:   []rune("aaa"),
			at:      0,
			want:    basicEditor{buf: []rune("aaa"), pos: 3},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 2},
			input:   []rune("aaa"),
			at:      -1,
			want:    basicEditor{buf: []rune("aaaAAA"), pos: 5},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 1},
			input:   []rune("aaa"),
			at:      2,
			want:    basicEditor{buf: []rune("AAaaaA"), pos: 1},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 3},
			input:   []rune("aaa"),
			at:      10,
			want:    basicEditor{buf: []rune("AAAaaa"), pos: 6},
		},
		{
			initial: basicEditor{buf: []rune("ABC"), pos: 0},
			input:   []rune("defg"),
			at:      0,
			want:    basicEditor{buf: []rune("defgABC"), pos: 3},
		},
	}
	for _, test := range tests {
		test.initial.insert(test.input, test.at)
		if !reflect.DeepEqual(test.initial, test.want) {
			t.Errorf("insert(%v, %v): got %v, want %v", test.input, test.at, test.initial, test.want)
		}
	}
}

func TestDelete(t *testing.T) {
	tests := []struct {
		initial basicEditor
		from    int
		to      int
		want    basicEditor
	}{
		{
			initial: basicEditor{buf: []rune(""), pos: 0},
			from:    -1,
			to:      1,
			want:    basicEditor{buf: []rune(""), pos: 0},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 2},
			from:    1,
			to:      2,
			want:    basicEditor{buf: []rune("AA"), pos: 1},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 1},
			from:    2,
			to:      0,
			want:    basicEditor{buf: []rune("A"), pos: 0},
		},
		{
			initial: basicEditor{buf: []rune("AAAABBCCaaaa"), pos: 3},
			from:    4,
			to:      8,
			want:    basicEditor{buf: []rune("AAAAaaaa"), pos: 3},
		},
		{
			initial: basicEditor{buf: []rune(""), pos: 0},
			from:    -1,
			to:      -1,
			want:    basicEditor{buf: []rune(""), pos: 0},
		},
	}
	for _, test := range tests {
		test.initial.delete(test.from, test.to)
		if !reflect.DeepEqual(test.initial, test.want) {
			t.Errorf("delete(%v, %v): got %v, want %v", test.from, test.to, test.initial, test.want)
		}
	}
}

func TestSlice(t *testing.T) {
	tests := []struct {
		initial []rune
		from    int
		to      int
		want    []rune
	}{
		{
			initial: []rune(""),
			from:    -1,
			to:      1,
			want:    []rune(""),
		},
		{
			initial: []rune("ABC"),
			from:    -1,
			to:      1,
			want:    []rune("A"),
		},
		{
			initial: []rune("aaa bbb ccc"),
			from:    2,
			to:      9,
			want:    []rune("a bbb c"),
		},
		{
			initial: []rune("aaa x bbb"),
			from:    10,
			to:      4,
			want:    []rune("x bbb"),
		},
	}
	for _, test := range tests {
		e := &basicEditor{buf: test.initial}
		got := e.slice(test.from, test.to)
		if string(got) != string(test.want) {
			t.Errorf("slice(%v, %v): got %v, want %v", test.from, test.to, got, test.want)
		}
	}
}

func TestIndex(t *testing.T) {
	tests := []struct {
		initial []rune
		ch      rune
		start   int
		want    int
	}{
		{
			initial: []rune(""),
			ch:      'a',
			start:   0,
			want:    -1,
		},
		{
			initial: []rune("aaa"),
			ch:      'a',
			start:   1,
			want:    1,
		},
		{
			initial: []rune("abcde"),
			ch:      'e',
			start:   3,
			want:    4,
		},
		{
			initial: []rune("AA AA"),
			ch:      'A',
			start:   2,
			want:    3,
		},
	}
	for _, test := range tests {
		e := &basicEditor{buf: test.initial}
		got := e.index(test.ch, test.start)
		if got != test.want {
			t.Errorf("index(%v, %v): got %v, want %v", test.ch, test.start, got, test.want)
		}
	}
}

func TestLastIndex(t *testing.T) {
	tests := []struct {
		initial []rune
		ch      rune
		start   int
		want    int
	}{
		{
			initial: []rune(""),
			ch:      'a',
			start:   0,
			want:    -1,
		},
		{
			initial: []rune("aaa"),
			ch:      'a',
			start:   1,
			want:    0,
		},
		{
			initial: []rune("abcde"),
			ch:      'e',
			start:   3,
			want:    -1,
		},
		{
			initial: []rune("AA AA"),
			ch:      'A',
			start:   2,
			want:    1,
		},
	}
	for _, test := range tests {
		e := &basicEditor{buf: test.initial}
		got := e.lastIndex(test.ch, test.start)
		if got != test.want {
			t.Errorf("lastIndex(%v, %v): got %v, want %v", string(test.ch), test.start, got, test.want)
		}
	}
}

func TestReplace(t *testing.T) {
	tests := []struct {
		initial basicEditor
		input   []rune
		at      int
		want    basicEditor
	}{
		{
			initial: basicEditor{buf: []rune(""), pos: 0},
			input:   []rune("aaa"),
			at:      0,
			want:    basicEditor{buf: []rune("aaa"), pos: 0},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 2},
			input:   []rune("aaa"),
			at:      -1,
			want:    basicEditor{buf: []rune("aaA"), pos: 2},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 1},
			input:   []rune("aaa"),
			at:      2,
			want:    basicEditor{buf: []rune("AAaaa"), pos: 1},
		},
		{
			initial: basicEditor{buf: []rune("AAA"), pos: 3},
			input:   []rune("aaa"),
			at:      10,
			want:    basicEditor{buf: []rune("AAA       aaa"), pos: 3},
		},
	}
	for _, test := range tests {
		test.initial.replace(test.input, test.at)
		if !reflect.DeepEqual(test.initial, test.want) {
			t.Errorf("replace(%v, %v): got %v, want %v", string(test.input), test.at, test.initial, test.want)
		}
	}
}
<|endoftext|>


<commit_before>package collector

import (
	"testing"

	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"github.com/smartystreets/goconvey/convey"
	"gopkg.in/DATA-DOG/go-sqlmock.v1"
)

func Test_ScrapeGlobalStatus(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatalf("error opening a stub database connection: %s", err)
	}
	defer db.Close()

	columns := []string{"Variable_name", "Value"}
	rows := sqlmock.NewRows(columns).
		AddRow("Com_alter_db", "1").
		AddRow("Com_show_status", "2").
		AddRow("Com_select", "3").
		AddRow("Connection_errors_internal", "4").
		AddRow("Handler_commit", "5").
		AddRow("Innodb_buffer_pool_pages_data", "6").
		AddRow("Innodb_buffer_pool_pages_flushed", "7").
		AddRow("Innodb_rows_read", "8").
		AddRow("Performance_schema_users_lost", "9").
		AddRow("Slave_running", "OFF").
		AddRow("Ssl_version", "").
		AddRow("Uptime", "10")
	mock.ExpectQuery(sanitizeQuery(globalStatusQuery)).WillReturnRows(rows)

	ch := make(chan prometheus.Metric)
	go func() {
		if err = ScrapeGlobalStatus(db, ch); err != nil {
			t.Errorf("error calling function on test: %s", err)
		}
		close(ch)
	}()

	counterExpected := []MetricResult{
		{labels: labelMap{"command": "alter_db"}, value: 1, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"command": "show_status"}, value: 2, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"command": "select"}, value: 3, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"error": "internal"}, value: 4, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"handler": "commit"}, value: 5, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"state": "data"}, value: 6, metricType: dto.MetricType_GAUGE},
		{labels: labelMap{"operation": "flushed"}, value: 7, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"operation": "read"}, value: 8, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"instrumentation": "users_lost"}, value: 9, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{}, value: 0, metricType: dto.MetricType_UNTYPED},
		{labels: labelMap{}, value: 10, metricType: dto.MetricType_UNTYPED},
	}
	convey.Convey("Metrics comparison", t, func() {
		for _, expect := range counterExpected {
			got := readMetric(<-ch)
			convey.So(got, convey.ShouldResemble, expect)
		}
	})

	// Ensure all SQL queries were executed
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Errorf("there were unfulfilled expections: %s", err)
	}
}
<commit_msg>Fix global_status test to use Go idiomatic name.<commit_after>package collector

import (
	"testing"

	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"github.com/smartystreets/goconvey/convey"
	"gopkg.in/DATA-DOG/go-sqlmock.v1"
)

func TestScrapeGlobalStatus(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatalf("error opening a stub database connection: %s", err)
	}
	defer db.Close()

	columns := []string{"Variable_name", "Value"}
	rows := sqlmock.NewRows(columns).
		AddRow("Com_alter_db", "1").
		AddRow("Com_show_status", "2").
		AddRow("Com_select", "3").
		AddRow("Connection_errors_internal", "4").
		AddRow("Handler_commit", "5").
		AddRow("Innodb_buffer_pool_pages_data", "6").
		AddRow("Innodb_buffer_pool_pages_flushed", "7").
		AddRow("Innodb_rows_read", "8").
		AddRow("Performance_schema_users_lost", "9").
		AddRow("Slave_running", "OFF").
		AddRow("Ssl_version", "").
		AddRow("Uptime", "10")
	mock.ExpectQuery(sanitizeQuery(globalStatusQuery)).WillReturnRows(rows)

	ch := make(chan prometheus.Metric)
	go func() {
		if err = ScrapeGlobalStatus(db, ch); err != nil {
			t.Errorf("error calling function on test: %s", err)
		}
		close(ch)
	}()

	counterExpected := []MetricResult{
		{labels: labelMap{"command": "alter_db"}, value: 1, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"command": "show_status"}, value: 2, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"command": "select"}, value: 3, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"error": "internal"}, value: 4, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"handler": "commit"}, value: 5, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"state": "data"}, value: 6, metricType: dto.MetricType_GAUGE},
		{labels: labelMap{"operation": "flushed"}, value: 7, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"operation": "read"}, value: 8, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{"instrumentation": "users_lost"}, value: 9, metricType: dto.MetricType_COUNTER},
		{labels: labelMap{}, value: 0, metricType: dto.MetricType_UNTYPED},
		{labels: labelMap{}, value: 10, metricType: dto.MetricType_UNTYPED},
	}
	convey.Convey("Metrics comparison", t, func() {
		for _, expect := range counterExpected {
			got := readMetric(<-ch)
			convey.So(got, convey.ShouldResemble, expect)
		}
	})

	// Ensure all SQL queries were executed
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Errorf("there were unfulfilled expections: %s", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/golang/protobuf/proto"
	"github.com/hashicorp/consul/api"
	"github.com/jeromer/syslogparser/rfc3164"
	"github.com/nsqio/go-nsq"
	"log"
	"net"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"
)

type StreamServer struct {
	*Setting
	exitChan      chan int
	msgChan       chan *LogFormat
	CurrentConfig map[string][]*regexp.Regexp
	wg            sync.WaitGroup
	client        *api.Client
	sync.Mutex
}

func (s *StreamServer) Run() {
	cfg := nsq.NewConfig()
	hostname, _ := os.Hostname()
	cfg.Set("user_agent", fmt.Sprintf("netstream/%s", hostname))
	cfg.Set("snappy", true)
	for i := 0; i < s.WritePoolSize; i++ {
		w, _ := nsq.NewProducer(s.NsqdAddr, cfg)
		go s.writeLoop(w)
	}
	ticker := time.Tick(time.Second * 600)
	go s.readUDP()
	go s.readTCP()
	var err error
	config := api.DefaultConfig()
	config.Address = s.ConsulAddress
	config.Datacenter = s.Datacenter
	config.Token = s.Token
	s.client, err = api.NewClient(config)
	if err != nil {
		fmt.Println("reload consul setting failed", err)
	}
	s.CurrentConfig, err = s.GetRegexp()
	for {
		select {
		case <-ticker:
			s.Lock()
			s.CurrentConfig, err = s.GetRegexp()
			s.Unlock()
			if err != nil {
				fmt.Println("reload consul setting failed", err)
			}
		case <-s.exitChan:
			return
		}
	}
}

func (s *StreamServer) writeLoop(w *nsq.Producer) {
	var bodies [][]byte
	for {
		select {
		case msg := <-s.msgChan:
			if s.IsIgnoreLog([]byte(msg.GetRawmsg())) {
				break
			}
			record, err := proto.Marshal(msg)
			if err != nil {
				break
			}
			bodies = append(bodies, record)
			if len(bodies) > 20 {
				w.MultiPublish(s.Topic, bodies)
				bodies = bodies[:0]
			}
		case <-s.exitChan:
			return
		}
	}
}
func (s *StreamServer) Stop() {
	close(s.exitChan)
	s.wg.Wait()
}

func (s *StreamServer) readUDP() {
	udpAddr, err := net.ResolveUDPAddr("udp", s.UdpPort)
	if err != nil {
		log.Fatal("udp:", err)
	}
	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		log.Fatal("server bind failed:", err)
	}
	defer conn.Close()
	buf := make([]byte, 8192*8)
	for {
		select {
		case <-s.exitChan:
			return
		default:
			size, addr, err := conn.ReadFromUDP(buf)
			if err != nil {
				log.Println("read log failed", err)
				continue
			}
			logFormat := &LogFormat{
				From:   proto.String(addr.String()),
				Rawmsg: proto.String(string(buf[:size])),
			}
			s.msgChan <- logFormat
		}
	}
}
func (s *StreamServer) readTCP() {
	server, err := net.Listen("tcp", s.TcpPort)
	if err != nil {
		log.Fatal("server bind failed:", err)
	}
	defer server.Close()
	for {
		select {
		case <-s.exitChan:
			return
		default:
			fd, err := server.Accept()
			if err != nil {
				log.Fatal("accept error", err)
				time.Sleep(time.Second)
			} else {
				go s.loghandle(fd)
			}
		}
	}
}

// receive log from tcp socket, encode json and send to msg_chan
func (s *StreamServer) loghandle(fd net.Conn) {
	defer fd.Close()
	scanner := bufio.NewScanner(fd)
	scanner.Split(bufio.ScanLines)
	addr := fd.RemoteAddr()
	s.wg.Add(1)
	defer s.wg.Done()
	var err error
	for {
		select {
		case <-s.exitChan:
			return
		default:
			if scanner.Scan() == false {
				err = scanner.Err()
				if err == nil {
					return
				}
			}
			if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
				log.Println(err)
				return
			}
			line := scanner.Text()
			logFormat := &LogFormat{
				From:   proto.String(addr.String()),
				Rawmsg: proto.String(string(line)),
			}
			s.msgChan <- logFormat
		}
	}
}

func (s *StreamServer) IsIgnoreLog(buf []byte) bool {
	if len(buf) < 1 {
		return true
	}
	p := rfc3164.NewParser(buf)
	if err := p.Parse(); err != nil {
		return false
	}
	data := p.Dump()
	tag := data["tag"].(string)
	if len(tag) == 0 {
		return false
	}
	s.Lock()
	rgs, ok := s.CurrentConfig[tag]
	s.Unlock()
	if ok {
		for _, r := range rgs {
			if r.MatchString(data["content"].(string)) {
				return true
			}
		}
	}
	return false
}
func (s *StreamServer) GetRegexp() (map[string][]*regexp.Regexp, error) {
	consulSetting := make(map[string][]*regexp.Regexp)
	kv := s.client.KV()
	pairs, _, err := kv.List(s.ConsulKey, nil)
	if err != nil {
		return consulSetting, err
	}
	size := len(s.ConsulKey) + 1
	for _, value := range pairs {
		if len(value.Key) > size {
			var regs []string
			if err := json.Unmarshal(value.Value, &regs); err == nil {
				var rs []*regexp.Regexp
				for _, v := range regs {
					x, e := regexp.CompilePOSIX(v)
					if e != nil {
						log.Println("get regexp", e)
						continue
					}
					rs = append(rs, x)
				}
				consulSetting[value.Key[size:]] = rs
			}
		}
	}
	return consulSetting, err
}
<commit_msg>improve batch size<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/golang/protobuf/proto"
	"github.com/hashicorp/consul/api"
	"github.com/jeromer/syslogparser/rfc3164"
	"github.com/nsqio/go-nsq"
	"log"
	"net"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"
)

type StreamServer struct {
	*Setting
	exitChan      chan int
	msgChan       chan *LogFormat
	CurrentConfig map[string][]*regexp.Regexp
	wg            sync.WaitGroup
	client        *api.Client
	sync.Mutex
}

func (s *StreamServer) Run() {
	cfg := nsq.NewConfig()
	hostname, _ := os.Hostname()
	cfg.Set("user_agent", fmt.Sprintf("netstream/%s", hostname))
	cfg.Set("snappy", true)
	for i := 0; i < s.WritePoolSize; i++ {
		w, _ := nsq.NewProducer(s.NsqdAddr, cfg)
		go s.writeLoop(w)
	}
	ticker := time.Tick(time.Second * 600)
	go s.readUDP()
	go s.readTCP()
	var err error
	config := api.DefaultConfig()
	config.Address = s.ConsulAddress
	config.Datacenter = s.Datacenter
	config.Token = s.Token
	s.client, err = api.NewClient(config)
	if err != nil {
		fmt.Println("reload consul setting failed", err)
	}
	s.CurrentConfig, err = s.GetRegexp()
	for {
		select {
		case <-ticker:
			s.Lock()
			s.CurrentConfig, err = s.GetRegexp()
			s.Unlock()
			if err != nil {
				fmt.Println("reload consul setting failed", err)
			}
		case <-s.exitChan:
			return
		}
	}
}

func (s *StreamServer) writeLoop(w *nsq.Producer) {
	var bodies [][]byte
	for {
		select {
		case msg := <-s.msgChan:
			if s.IsIgnoreLog([]byte(msg.GetRawmsg())) {
				break
			}
			record, err := proto.Marshal(msg)
			if err != nil {
				break
			}
			bodies = append(bodies, record)
			if len(bodies) > 100 {
				w.MultiPublish(s.Topic, bodies)
				bodies = bodies[:0]
			}
		case <-s.exitChan:
			return
		}
	}
}
func (s *StreamServer) Stop() {
	close(s.exitChan)
	s.wg.Wait()
}

func (s *StreamServer) readUDP() {
	udpAddr, err := net.ResolveUDPAddr("udp", s.UdpPort)
	if err != nil {
		log.Fatal("udp:", err)
	}
	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		log.Fatal("server bind failed:", err)
	}
	defer conn.Close()
	buf := make([]byte, 8192*8)
	for {
		select {
		case <-s.exitChan:
			return
		default:
			size, addr, err := conn.ReadFromUDP(buf)
			if err != nil {
				log.Println("read log failed", err)
				continue
			}
			logFormat := &LogFormat{
				From:   proto.String(addr.String()),
				Rawmsg: proto.String(string(buf[:size])),
			}
			s.msgChan <- logFormat
		}
	}
}
func (s *StreamServer) readTCP() {
	server, err := net.Listen("tcp", s.TcpPort)
	if err != nil {
		log.Fatal("server bind failed:", err)
	}
	defer server.Close()
	for {
		select {
		case <-s.exitChan:
			return
		default:
			fd, err := server.Accept()
			if err != nil {
				log.Fatal("accept error", err)
				time.Sleep(time.Second)
			} else {
				go s.loghandle(fd)
			}
		}
	}
}

// receive log from tcp socket, encode json and send to msg_chan
func (s *StreamServer) loghandle(fd net.Conn) {
	defer fd.Close()
	scanner := bufio.NewScanner(fd)
	scanner.Split(bufio.ScanLines)
	addr := fd.RemoteAddr()
	s.wg.Add(1)
	defer s.wg.Done()
	var err error
	for {
		select {
		case <-s.exitChan:
			return
		default:
			if scanner.Scan() == false {
				err = scanner.Err()
				if err == nil {
					return
				}
			}
			if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
				log.Println(err)
				return
			}
			line := scanner.Text()
			logFormat := &LogFormat{
				From:   proto.String(addr.String()),
				Rawmsg: proto.String(string(line)),
			}
			s.msgChan <- logFormat
		}
	}
}

func (s *StreamServer) IsIgnoreLog(buf []byte) bool {
	if len(buf) < 1 {
		return true
	}
	p := rfc3164.NewParser(buf)
	if err := p.Parse(); err != nil {
		return false
	}
	data := p.Dump()
	tag := data["tag"].(string)
	if len(tag) == 0 {
		return false
	}
	s.Lock()
	rgs, ok := s.CurrentConfig[tag]
	s.Unlock()
	if ok {
		for _, r := range rgs {
			if r.MatchString(data["content"].(string)) {
				return true
			}
		}
	}
	return false
}
func (s *StreamServer) GetRegexp() (map[string][]*regexp.Regexp, error) {
	consulSetting := make(map[string][]*regexp.Regexp)
	kv := s.client.KV()
	pairs, _, err := kv.List(s.ConsulKey, nil)
	if err != nil {
		return consulSetting, err
	}
	size := len(s.ConsulKey) + 1
	for _, value := range pairs {
		if len(value.Key) > size {
			var regs []string
			if err := json.Unmarshal(value.Value, &regs); err == nil {
				var rs []*regexp.Regexp
				for _, v := range regs {
					x, e := regexp.CompilePOSIX(v)
					if e != nil {
						log.Println("get regexp", e)
						continue
					}
					rs = append(rs, x)
				}
				consulSetting[value.Key[size:]] = rs
			}
		}
	}
	return consulSetting, err
}
<|endoftext|>


<commit_before><commit_msg>Remove unused constant, follow up from 111089f (#712)<commit_after><|endoftext|>


<commit_before><commit_msg>update osd_test<commit_after><|endoftext|>


<commit_before>package heppdt_test

import (
	"testing"

	"github.com/go-hep/heppdt"
)

func TestPID(t *testing.T) {
	/*
	ids := []int{ 5, 25, 15, 213, -3214, 10213, 9050225, -200543, 129050225,
		2000025, 3101, 3301, -2212, 1000020040, -1000060120, 555,
		5000040, 5100005, 24, 5100024, 5100025, 9221132,
		4111370, -4120240, 4110050, 10013730,
		1000993, 1000612, 1000622, 1000632, 1006213, 1000652,
		1009113, 1009213, 1009323,
		1093114, 1009333, 1006313, 1092214, 1006223,
	}
	*/

	for _, table := range []struct{
		id int
		nx, nr, n1, nq1, nq2, nq3, nj int
		extra int
		jspin, lspin, spin int
		fid int
		charge float64
		valid bool

		ishadron, ismeson,isrhadron,isbaryon bool
		islepton, isdiquark, isnucleus bool

		issusy, isdyon, isqlb, isunknown bool

		hasup, hasdown, hasstrange, hascharm, hasbottom, hastop bool
	}{
		{
			id: 5,
			nx: 0, nr:0, n1:0, nq1:0, nq2:0, nq3:0, nj:5,
			extra: 0,
			jspin: 2, lspin:0, spin:0,
			fid: 0,
			charge: -1./3.,
			valid: true,
		},
	} {
		id := table.id
		pid := heppdt.PID(id)
		nx := pid.Digit(heppdt.N)
		if nx != table.nx {
			t.Fatalf("pid=%d. expected nx=%d. got=%d", id, table.nx, nx)
		}

		nr := pid.Digit(heppdt.Nr)
		if nr != table.nr {
			t.Fatalf("pid=%d. expected nr=%d. got=%d", id, table.nr, nr)
		}

		nq1 := pid.Digit(heppdt.Nq1)
		if nq1 != table.nq1 {
			t.Fatalf("pid=%d. expected nq1=%d. got=%d", id, table.nq1, nq1)
		}

		nq2 := pid.Digit(heppdt.Nq2)
		if nq2 != table.nq2 {
			t.Fatalf("pid=%d. expected nq2=%d. got=%d", id, table.nq2, nq2)
		}

		nq3 := pid.Digit(heppdt.Nq3)
		if nq3 != table.nq3 {
			t.Fatalf("pid=%d. expected nq3=%d. got=%d", id, table.nq3, nq3)
		}

		extra := pid.ExtraBits()
		if extra != table.extra {
			t.Fatalf("pid=%d. expected extra=%d. got=%d", id, table.extra, extra)
		}

		jspin := pid.JSpin()
		if jspin != table.jspin {
			t.Fatalf("pid=%d. expected jspin=%d. got=%d", id, table.jspin, jspin)
		}

		lspin := pid.LSpin()
		if lspin != table.lspin {
			t.Fatalf("pid=%d. expected lspin=%d. got=%d", id, table.lspin, lspin)
		}

		spin := pid.SSpin()
		if spin != table.spin {
			t.Fatalf("pid=%d. expected sspin=%d. got=%d", id, table.spin, spin)
		}

		charge := pid.Charge()
		if charge != table.charge {
			t.Fatalf("pid=%d. expected charge=%v. got=%v", id, table.charge, charge)
		}

		if pid.IsValid() != table.valid {
			t.Fatalf("expected pid=%d valid=%v. got=%v", int(pid), table.valid, pid.IsValid())
		}

		if table.ishadron != pid.IsHadron() {
			t.Fatalf("expected pid=%d ishadron=%v. got=%v", int(pid), table.ishadron, pid.IsHadron())
		}

		if table.ismeson != pid.IsMeson() {
			t.Fatalf("expected pid=%d ismeson=%v. got=%v", int(pid), table.ismeson, pid.IsMeson())
		}

		if table.isrhadron != pid.IsRhadron() {
			t.Fatalf("expected pid=%d isrhadron=%v. got=%v", int(pid), table.isrhadron, pid.IsRhadron())
		}

		if table.isbaryon != pid.IsBaryon() {
			t.Fatalf("expected pid=%d isbaryon=%v. got=%v", int(pid), table.isbaryon, pid.IsBaryon())
		}

		if table.islepton != pid.IsLepton() {
			t.Fatalf("expected pid=%d islepton=%v. got=%v", int(pid), table.islepton, pid.IsLepton())
		}

		if table.isdiquark != pid.IsDiQuark() {
			t.Fatalf("expected pid=%d isdiquark=%v. got=%v", int(pid), table.isdiquark, pid.IsDiQuark())
		}

		if table.isnucleus != pid.IsNucleus() {
			t.Fatalf("expected pid=%d isnucleus=%v. got=%v", int(pid), table.isnucleus, pid.IsNucleus())
		}

		if table.issusy != pid.IsSUSY() {
			t.Fatalf("expected pid=%d issusy=%v. got=%v", int(pid), table.issusy, pid.IsSUSY())
		}

		if table.isdyon != pid.IsDyon() {
			t.Fatalf("expected pid=%d isdyon=%v. got=%v", int(pid), table.isdyon, pid.IsDyon())
		}

		if table.isqlb != (pid.AbsPID() < 100) {
			t.Fatalf("expected pid=%d is-quark/lepton/boson=%v. got=%v", int(pid), table.isqlb, pid.AbsPID() < 100)
		}

		if table.isunknown != (!pid.IsSUSY() && !pid.IsDyon() && !pid.IsHadron() &&
			!pid.IsLepton()	&& !pid.IsDiQuark() && !pid.IsNucleus() &&
			!(pid.AbsPID() < 100)) {
			t.Fatalf("expected pid=%d isunknown=%v. got=%v", int(pid), table.isunknown, !table.isunknown)
		}

		if table.hasup != pid.HasUp() {
			t.Fatalf("expected pid=%d has=%v. got=%v", int(pid), table.hasup, pid.HasUp())
		}

		if table.hasdown != pid.HasDown() {
			t.Fatalf("expected pid=%d hasdown=%v. got=%v", int(pid), table.hasdown, pid.HasDown())
		}

		if table.hasstrange != pid.HasStrange() {
			t.Fatalf("expected pid=%d hasstrange=%v. got=%v", int(pid), table.hasstrange, pid.HasStrange())
		}

		if table.hascharm != pid.HasCharm() {
			t.Fatalf("expected pid=%d hascharm=%v. got=%v", int(pid), table.hascharm, pid.HasCharm())
		}

		if table.hasbottom != pid.HasBottom() {
			t.Fatalf("expected pid=%d hasbottom=%v. got=%v", int(pid), table.hasbottom, pid.HasBottom())
		}

		if table.hastop != pid.HasTop() {
			t.Fatalf("expected pid=%d hastop=%v. got=%v", int(pid), table.hastop, pid.HasTop())
		}
	}
}
<commit_msg>test: complete definition for test-5<commit_after>package heppdt_test

import (
	"testing"

	"github.com/go-hep/heppdt"
)

func TestPID(t *testing.T) {
	/*
	ids := []int{ 5, 25, 15, 213, -3214, 10213, 9050225, -200543, 129050225,
		2000025, 3101, 3301, -2212, 1000020040, -1000060120, 555,
		5000040, 5100005, 24, 5100024, 5100025, 9221132,
		4111370, -4120240, 4110050, 10013730,
		1000993, 1000612, 1000622, 1000632, 1006213, 1000652,
		1009113, 1009213, 1009323,
		1093114, 1009333, 1006313, 1092214, 1006223,
	}
	*/

	for _, table := range []struct{
		id int
		nx, nr, n1, nq1, nq2, nq3, nj int
		extra int
		jspin, lspin, spin int
		fid int
		charge float64
		valid bool

		ishadron, ismeson,isrhadron,isbaryon bool
		islepton, isdiquark, isnucleus bool

		issusy, isdyon, isqlb, isunknown bool

		hasup, hasdown, hasstrange, hascharm, hasbottom, hastop bool
	}{
		{
			id: 5,
			nx: 0, nr:0, n1:0, nq1:0, nq2:0, nq3:0, nj:5,
			extra: 0,
			jspin: 2, lspin:0, spin:0,
			fid: 0,
			charge: -1./3.,
			valid: true,
			isqlb: true,
		},
	} {
		id := table.id
		pid := heppdt.PID(id)
		nx := pid.Digit(heppdt.N)
		if nx != table.nx {
			t.Fatalf("pid=%d. expected nx=%d. got=%d", id, table.nx, nx)
		}

		nr := pid.Digit(heppdt.Nr)
		if nr != table.nr {
			t.Fatalf("pid=%d. expected nr=%d. got=%d", id, table.nr, nr)
		}

		nq1 := pid.Digit(heppdt.Nq1)
		if nq1 != table.nq1 {
			t.Fatalf("pid=%d. expected nq1=%d. got=%d", id, table.nq1, nq1)
		}

		nq2 := pid.Digit(heppdt.Nq2)
		if nq2 != table.nq2 {
			t.Fatalf("pid=%d. expected nq2=%d. got=%d", id, table.nq2, nq2)
		}

		nq3 := pid.Digit(heppdt.Nq3)
		if nq3 != table.nq3 {
			t.Fatalf("pid=%d. expected nq3=%d. got=%d", id, table.nq3, nq3)
		}

		extra := pid.ExtraBits()
		if extra != table.extra {
			t.Fatalf("pid=%d. expected extra=%d. got=%d", id, table.extra, extra)
		}

		jspin := pid.JSpin()
		if jspin != table.jspin {
			t.Fatalf("pid=%d. expected jspin=%d. got=%d", id, table.jspin, jspin)
		}

		lspin := pid.LSpin()
		if lspin != table.lspin {
			t.Fatalf("pid=%d. expected lspin=%d. got=%d", id, table.lspin, lspin)
		}

		spin := pid.SSpin()
		if spin != table.spin {
			t.Fatalf("pid=%d. expected sspin=%d. got=%d", id, table.spin, spin)
		}

		charge := pid.Charge()
		if charge != table.charge {
			t.Fatalf("pid=%d. expected charge=%v. got=%v", id, table.charge, charge)
		}

		if pid.IsValid() != table.valid {
			t.Fatalf("expected pid=%d valid=%v. got=%v", int(pid), table.valid, pid.IsValid())
		}

		if table.ishadron != pid.IsHadron() {
			t.Fatalf("expected pid=%d ishadron=%v. got=%v", int(pid), table.ishadron, pid.IsHadron())
		}

		if table.ismeson != pid.IsMeson() {
			t.Fatalf("expected pid=%d ismeson=%v. got=%v", int(pid), table.ismeson, pid.IsMeson())
		}

		if table.isrhadron != pid.IsRhadron() {
			t.Fatalf("expected pid=%d isrhadron=%v. got=%v", int(pid), table.isrhadron, pid.IsRhadron())
		}

		if table.isbaryon != pid.IsBaryon() {
			t.Fatalf("expected pid=%d isbaryon=%v. got=%v", int(pid), table.isbaryon, pid.IsBaryon())
		}

		if table.islepton != pid.IsLepton() {
			t.Fatalf("expected pid=%d islepton=%v. got=%v", int(pid), table.islepton, pid.IsLepton())
		}

		if table.isdiquark != pid.IsDiQuark() {
			t.Fatalf("expected pid=%d isdiquark=%v. got=%v", int(pid), table.isdiquark, pid.IsDiQuark())
		}

		if table.isnucleus != pid.IsNucleus() {
			t.Fatalf("expected pid=%d isnucleus=%v. got=%v", int(pid), table.isnucleus, pid.IsNucleus())
		}

		if table.issusy != pid.IsSUSY() {
			t.Fatalf("expected pid=%d issusy=%v. got=%v", int(pid), table.issusy, pid.IsSUSY())
		}

		if table.isdyon != pid.IsDyon() {
			t.Fatalf("expected pid=%d isdyon=%v. got=%v", int(pid), table.isdyon, pid.IsDyon())
		}

		if table.isqlb != (pid.AbsPID() < 100) {
			t.Fatalf("expected pid=%d is-quark/lepton/boson=%v. got=%v", int(pid), table.isqlb, pid.AbsPID() < 100)
		}

		if table.isunknown != (!pid.IsSUSY() && !pid.IsDyon() && !pid.IsHadron() &&
			!pid.IsLepton()	&& !pid.IsDiQuark() && !pid.IsNucleus() &&
			!(pid.AbsPID() < 100)) {
			t.Fatalf("expected pid=%d isunknown=%v. got=%v", int(pid), table.isunknown, !table.isunknown)
		}

		if table.hasup != pid.HasUp() {
			t.Fatalf("expected pid=%d has=%v. got=%v", int(pid), table.hasup, pid.HasUp())
		}

		if table.hasdown != pid.HasDown() {
			t.Fatalf("expected pid=%d hasdown=%v. got=%v", int(pid), table.hasdown, pid.HasDown())
		}

		if table.hasstrange != pid.HasStrange() {
			t.Fatalf("expected pid=%d hasstrange=%v. got=%v", int(pid), table.hasstrange, pid.HasStrange())
		}

		if table.hascharm != pid.HasCharm() {
			t.Fatalf("expected pid=%d hascharm=%v. got=%v", int(pid), table.hascharm, pid.HasCharm())
		}

		if table.hasbottom != pid.HasBottom() {
			t.Fatalf("expected pid=%d hasbottom=%v. got=%v", int(pid), table.hasbottom, pid.HasBottom())
		}

		if table.hastop != pid.HasTop() {
			t.Fatalf("expected pid=%d hastop=%v. got=%v", int(pid), table.hastop, pid.HasTop())
		}
	}
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package ciliumnode

import (
	"errors"
	"fmt"
	"os"
	"path"
	"reflect"
	"testing"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"

	operatorOption "github.com/cilium/cilium/operator/option"
	v2 "github.com/cilium/cilium/pkg/k8s/apis/cilium.io/v2"
	agentOption "github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/test/controlplane/suite"
)

type step struct {
	filename       string
	expectedLabels map[string]string
}

var steps = []step{
	{
		"state1.yaml",
		map[string]string{
			"beta.kubernetes.io/arch": "amd64",
			"beta.kubernetes.io/os":   "linux",
			"cilium.io/ci-node":       "k8s1",
			"kubernetes.io/arch":      "amd64",
			"kubernetes.io/hostname":  "cilium-nodes-worker",
			"kubernetes.io/os":        "linux",
			"test-label":              "test-value",
		},
	},

	{
		"state2.yaml",
		map[string]string{
			"beta.kubernetes.io/arch": "amd64",
			"beta.kubernetes.io/os":   "linux",
			"cilium.io/ci-node":       "k8s1",
			"kubernetes.io/arch":      "amd64",
			"kubernetes.io/hostname":  "cilium-nodes-worker",
			"kubernetes.io/os":        "linux",
		},
	},

	{
		"state3.yaml",
		map[string]string{
			"beta.kubernetes.io/arch": "amd64",
			"beta.kubernetes.io/os":   "linux",
			"cilium.io/ci-node":       "k8s1",
			"kubernetes.io/arch":      "amd64",
			"kubernetes.io/hostname":  "cilium-nodes-worker",
			"kubernetes.io/os":        "linux",

			"another-test-label": "another-test-value",
		},
	},

	{
		"state4.yaml",
		map[string]string{
			"beta.kubernetes.io/arch": "amd64",
			"beta.kubernetes.io/os":   "linux",
			"cilium.io/ci-node":       "k8s1",
			"kubernetes.io/arch":      "amd64",
			"kubernetes.io/hostname":  "cilium-nodes-worker",
			"kubernetes.io/os":        "linux",

			"another-test-label": "changed-test-value",
		},
	},
}

func init() {
	suite.AddTestCase("CiliumNodes", func(t *testing.T) {
		cwd, err := os.Getwd()
		if err != nil {
			t.Fatal(err)
		}

		modConfig := func(daemonCfg *agentOption.DaemonConfig, _ *operatorOption.OperatorConfig) {
			daemonCfg.EnableNodePort = true
		}
		for _, version := range []string{"1.20", "1.22", "1.24"} {
			abs := func(f string) string { return path.Join(cwd, "node", "ciliumnodes", "v"+version, f) }

			t.Run("v"+version, func(t *testing.T) {
				test := suite.NewControlPlaneTest(t, "cilium-nodes-control-plane", version)

				// Feed in initial state and start the agent.
				test.
					UpdateObjectsFromFile(abs("init.yaml")).
					SetupEnvironment(modConfig).
					StartAgent()

				// Run through the test steps
				for _, step := range steps {
					test.UpdateObjectsFromFile(abs(step.filename))
					test.Eventually(func() error { return validateLabels(test, step.expectedLabels) })
				}

				test.StopAgent()
			})
		}
	})
}

func validateLabels(test *suite.ControlPlaneTest, expectedLabels map[string]string) error {
	nodeLabels, err := getTestNodeLabels(test, "cilium-nodes-worker")
	if err != nil {
		return fmt.Errorf("failed to get Node labels: %w", err)
	}
	if !reflect.DeepEqual(expectedLabels, nodeLabels) {
		return fmt.Errorf("Node labels mismatch, expected: %v, found: %v",
			nodeLabels, expectedLabels)
	}
	ciliumNodeLabels, err := getTestCiliumNodeLabels(test, "cilium-nodes-worker")
	if err != nil {
		return fmt.Errorf("failed to get CiliumNode labels: %w", err)
	}
	if !reflect.DeepEqual(expectedLabels, ciliumNodeLabels) {
		return fmt.Errorf("CiliumNode labels mismatch, expected: %v, found: %v",
			ciliumNodeLabels, expectedLabels)
	}
	return nil
}

func getTestNodeLabels(test *suite.ControlPlaneTest, name string) (map[string]string, error) {
	nodeObj, err := test.Get(
		schema.GroupVersionResource{Group: "", Version: "v1", Resource: "nodes"},
		"",
		name,
	)
	if err != nil {
		return nil, fmt.Errorf("unable to get %q Node: %w", name, err)
	}
	node, ok := nodeObj.(*v1.Node)
	if !ok {
		return nil, errors.New("type assertion failed for Node obj")
	}

	return node.GetLabels(), nil
}

func getTestCiliumNodeLabels(test *suite.ControlPlaneTest, name string) (map[string]string, error) {
	ciliumNodeObj, err := test.Get(
		schema.GroupVersionResource{Group: "cilium.io", Version: "v2", Resource: "ciliumnodes"},
		"",
		name,
	)
	if err != nil {
		return nil, fmt.Errorf("unable to get %q CiliumNode: %w", name, err)
	}
	ciliumNode, ok := ciliumNodeObj.(*v2.CiliumNode)
	if !ok {
		return nil, errors.New("type assertion failed for CiliumNode obj")
	}

	return ciliumNode.GetLabels(), nil
}
<commit_msg>controlplane: use right node name in tests<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package ciliumnode

import (
	"errors"
	"fmt"
	"os"
	"path"
	"reflect"
	"testing"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"

	operatorOption "github.com/cilium/cilium/operator/option"
	v2 "github.com/cilium/cilium/pkg/k8s/apis/cilium.io/v2"
	agentOption "github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/test/controlplane/suite"
)

type step struct {
	filename       string
	expectedLabels map[string]string
}

var steps = []step{
	{
		"state1.yaml",
		map[string]string{
			"beta.kubernetes.io/arch": "amd64",
			"beta.kubernetes.io/os":   "linux",
			"cilium.io/ci-node":       "k8s1",
			"kubernetes.io/arch":      "amd64",
			"kubernetes.io/hostname":  "cilium-nodes-worker",
			"kubernetes.io/os":        "linux",
			"test-label":              "test-value",
		},
	},

	{
		"state2.yaml",
		map[string]string{
			"beta.kubernetes.io/arch": "amd64",
			"beta.kubernetes.io/os":   "linux",
			"cilium.io/ci-node":       "k8s1",
			"kubernetes.io/arch":      "amd64",
			"kubernetes.io/hostname":  "cilium-nodes-worker",
			"kubernetes.io/os":        "linux",
		},
	},

	{
		"state3.yaml",
		map[string]string{
			"beta.kubernetes.io/arch": "amd64",
			"beta.kubernetes.io/os":   "linux",
			"cilium.io/ci-node":       "k8s1",
			"kubernetes.io/arch":      "amd64",
			"kubernetes.io/hostname":  "cilium-nodes-worker",
			"kubernetes.io/os":        "linux",

			"another-test-label": "another-test-value",
		},
	},

	{
		"state4.yaml",
		map[string]string{
			"beta.kubernetes.io/arch": "amd64",
			"beta.kubernetes.io/os":   "linux",
			"cilium.io/ci-node":       "k8s1",
			"kubernetes.io/arch":      "amd64",
			"kubernetes.io/hostname":  "cilium-nodes-worker",
			"kubernetes.io/os":        "linux",

			"another-test-label": "changed-test-value",
		},
	},
}

func init() {
	suite.AddTestCase("CiliumNodes", func(t *testing.T) {
		cwd, err := os.Getwd()
		if err != nil {
			t.Fatal(err)
		}

		modConfig := func(daemonCfg *agentOption.DaemonConfig, _ *operatorOption.OperatorConfig) {
			daemonCfg.EnableNodePort = true
		}
		for _, version := range []string{"1.20", "1.22", "1.24"} {
			abs := func(f string) string { return path.Join(cwd, "node", "ciliumnodes", "v"+version, f) }

			t.Run("v"+version, func(t *testing.T) {
				test := suite.NewControlPlaneTest(t, "cilium-nodes-worker", version)

				// Feed in initial state and start the agent.
				test.
					UpdateObjectsFromFile(abs("init.yaml")).
					SetupEnvironment(modConfig).
					StartAgent()

				// Run through the test steps
				for _, step := range steps {
					test.UpdateObjectsFromFile(abs(step.filename))
					test.Eventually(func() error { return validateLabels(test, step.expectedLabels) })
				}

				test.StopAgent()
			})
		}
	})
}

func validateLabels(test *suite.ControlPlaneTest, expectedLabels map[string]string) error {
	nodeLabels, err := getTestNodeLabels(test, "cilium-nodes-worker")
	if err != nil {
		return fmt.Errorf("failed to get Node labels: %w", err)
	}
	if !reflect.DeepEqual(expectedLabels, nodeLabels) {
		return fmt.Errorf("Node labels mismatch, expected: %v, found: %v",
			nodeLabels, expectedLabels)
	}
	ciliumNodeLabels, err := getTestCiliumNodeLabels(test, "cilium-nodes-worker")
	if err != nil {
		return fmt.Errorf("failed to get CiliumNode labels: %w", err)
	}
	if !reflect.DeepEqual(expectedLabels, ciliumNodeLabels) {
		return fmt.Errorf("CiliumNode labels mismatch, expected: %v, found: %v",
			ciliumNodeLabels, expectedLabels)
	}
	return nil
}

func getTestNodeLabels(test *suite.ControlPlaneTest, name string) (map[string]string, error) {
	nodeObj, err := test.Get(
		schema.GroupVersionResource{Group: "", Version: "v1", Resource: "nodes"},
		"",
		name,
	)
	if err != nil {
		return nil, fmt.Errorf("unable to get %q Node: %w", name, err)
	}
	node, ok := nodeObj.(*v1.Node)
	if !ok {
		return nil, errors.New("type assertion failed for Node obj")
	}

	return node.GetLabels(), nil
}

func getTestCiliumNodeLabels(test *suite.ControlPlaneTest, name string) (map[string]string, error) {
	ciliumNodeObj, err := test.Get(
		schema.GroupVersionResource{Group: "cilium.io", Version: "v2", Resource: "ciliumnodes"},
		"",
		name,
	)
	if err != nil {
		return nil, fmt.Errorf("unable to get %q CiliumNode: %w", name, err)
	}
	ciliumNode, ok := ciliumNodeObj.(*v2.CiliumNode)
	if !ok {
		return nil, errors.New("type assertion failed for CiliumNode obj")
	}

	return ciliumNode.GetLabels(), nil
}
<|endoftext|>