<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
)

func resourceAwsKmsKey() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsKmsKeyCreate,
		Read:   resourceAwsKmsKeyRead,
		Update: resourceAwsKmsKeyUpdate,
		Delete: resourceAwsKmsKeyDelete,

		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"key_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"key_usage": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, es []error) {
					value := v.(string)
					if !(value == "ENCRYPT_DECRYPT" || value == "") {
						es = append(es, fmt.Errorf(
							"%q must be ENCRYPT_DECRYPT or not specified", k))
					}
					return
				},
			},
			"policy": &schema.Schema{
				Type:             schema.TypeString,
				Optional:         true,
				Computed:         true,
				ValidateFunc:     validateJsonString,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},
			"is_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"enable_key_rotation": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},
			"deletion_window_in_days": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, es []error) {
					value := v.(int)
					if value > 30 || value < 7 {
						es = append(es, fmt.Errorf(
							"%q must be between 7 and 30 days inclusive", k))
					}
					return
				},
			},
		},
	}
}

func resourceAwsKmsKeyCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn

	// Allow aws to chose default values if we don't pass them
	var req kms.CreateKeyInput
	if v, exists := d.GetOk("description"); exists {
		req.Description = aws.String(v.(string))
	}
	if v, exists := d.GetOk("key_usage"); exists {
		req.KeyUsage = aws.String(v.(string))
	}
	if v, exists := d.GetOk("policy"); exists {
		req.Policy = aws.String(v.(string))
	}

	resp, err := conn.CreateKey(&req)
	if err != nil {
		return err
	}

	d.SetId(*resp.KeyMetadata.KeyId)
	d.Set("key_id", resp.KeyMetadata.KeyId)

	return _resourceAwsKmsKeyUpdate(d, meta, true)
}

func resourceAwsKmsKeyRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn

	req := &kms.DescribeKeyInput{
		KeyId: aws.String(d.Id()),
	}
	resp, err := conn.DescribeKey(req)
	if err != nil {
		return err
	}
	metadata := resp.KeyMetadata

	if *metadata.KeyState == "PendingDeletion" {
		log.Printf("[WARN] Removing KMS key %s because it's already gone", d.Id())
		d.SetId("")
		return nil
	}

	d.SetId(*metadata.KeyId)

	d.Set("arn", metadata.Arn)
	d.Set("key_id", metadata.KeyId)
	d.Set("description", metadata.Description)
	d.Set("key_usage", metadata.KeyUsage)
	d.Set("is_enabled", metadata.Enabled)

	p, err := conn.GetKeyPolicy(&kms.GetKeyPolicyInput{
		KeyId:      metadata.KeyId,
		PolicyName: aws.String("default"),
	})
	if err != nil {
		return err
	}

	policy, _ := normalizeJsonString(*p.Policy)
	d.Set("policy", policy)

	krs, err := conn.GetKeyRotationStatus(&kms.GetKeyRotationStatusInput{
		KeyId: metadata.KeyId,
	})
	if err != nil {
		return err
	}
	d.Set("enable_key_rotation", krs.KeyRotationEnabled)

	return nil
}

func resourceAwsKmsKeyUpdate(d *schema.ResourceData, meta interface{}) error {
	return _resourceAwsKmsKeyUpdate(d, meta, false)
}

// We expect new keys to be enabled already
// but there is no easy way to differentiate between Update()
// called from Create() and regular update, so we have this wrapper
func _resourceAwsKmsKeyUpdate(d *schema.ResourceData, meta interface{}, isFresh bool) error {
	conn := meta.(*AWSClient).kmsconn

	if d.HasChange("is_enabled") && d.Get("is_enabled").(bool) && !isFresh {
		// Enable before any attributes will be modified
		if err := updateKmsKeyStatus(conn, d.Id(), d.Get("is_enabled").(bool)); err != nil {
			return err
		}
	}

	if d.HasChange("enable_key_rotation") {
		if err := updateKmsKeyRotationStatus(conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("description") {
		if err := resourceAwsKmsKeyDescriptionUpdate(conn, d); err != nil {
			return err
		}
	}
	if d.HasChange("policy") {
		if err := resourceAwsKmsKeyPolicyUpdate(conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("is_enabled") && !d.Get("is_enabled").(bool) {
		// Only disable when all attributes are modified
		// because we cannot modify disabled keys
		if err := updateKmsKeyStatus(conn, d.Id(), d.Get("is_enabled").(bool)); err != nil {
			return err
		}
	}

	return resourceAwsKmsKeyRead(d, meta)
}

func resourceAwsKmsKeyDescriptionUpdate(conn *kms.KMS, d *schema.ResourceData) error {
	description := d.Get("description").(string)
	keyId := d.Get("key_id").(string)

	log.Printf("[DEBUG] KMS key: %s, update description: %s", keyId, description)

	req := &kms.UpdateKeyDescriptionInput{
		Description: aws.String(description),
		KeyId:       aws.String(keyId),
	}
	_, err := conn.UpdateKeyDescription(req)
	return err
}

func resourceAwsKmsKeyPolicyUpdate(conn *kms.KMS, d *schema.ResourceData) error {
	policy, _ := normalizeJsonString(d.Get("policy").(string))
	keyId := d.Get("key_id").(string)

	log.Printf("[DEBUG] KMS key: %s, update policy: %s", keyId, policy)

	req := &kms.PutKeyPolicyInput{
		KeyId:      aws.String(keyId),
		Policy:     aws.String(policy),
		PolicyName: aws.String("default"),
	}
	_, err := conn.PutKeyPolicy(req)
	return err
}

func updateKmsKeyStatus(conn *kms.KMS, id string, shouldBeEnabled bool) error {
	var err error

	if shouldBeEnabled {
		log.Printf("[DEBUG] Enabling KMS key %q", id)
		_, err = conn.EnableKey(&kms.EnableKeyInput{
			KeyId: aws.String(id),
		})
	} else {
		log.Printf("[DEBUG] Disabling KMS key %q", id)
		_, err = conn.DisableKey(&kms.DisableKeyInput{
			KeyId: aws.String(id),
		})
	}

	if err != nil {
		return fmt.Errorf("Failed to set KMS key %q status to %t: %q",
			id, shouldBeEnabled, err.Error())
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{fmt.Sprintf("%t", !shouldBeEnabled)},
		Target:                    []string{fmt.Sprintf("%t", shouldBeEnabled)},
		Timeout:                   20 * time.Minute,
		MinTimeout:                2 * time.Second,
		ContinuousTargetOccurence: 10,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s enabled status is %t",
				id, shouldBeEnabled)
			resp, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(id),
			})
			if err != nil {
				return resp, "FAILED", err
			}
			status := fmt.Sprintf("%t", *resp.KeyMetadata.Enabled)
			log.Printf("[DEBUG] KMS key %s status received: %s, retrying", id, status)

			return resp, status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed setting KMS key status to %t: %s", shouldBeEnabled, err)
	}

	return nil
}

func updateKmsKeyRotationStatus(conn *kms.KMS, d *schema.ResourceData) error {
	var err error
	shouldEnableRotation := d.Get("enable_key_rotation").(bool)
	if shouldEnableRotation {
		log.Printf("[DEBUG] Enabling key rotation for KMS key %q", d.Id())
		_, err = conn.EnableKeyRotation(&kms.EnableKeyRotationInput{
			KeyId: aws.String(d.Id()),
		})
	} else {
		log.Printf("[DEBUG] Disabling key rotation for KMS key %q", d.Id())
		_, err = conn.DisableKeyRotation(&kms.DisableKeyRotationInput{
			KeyId: aws.String(d.Id()),
		})
	}

	if err != nil {
		return fmt.Errorf("Failed to set key rotation for %q to %t: %q",
			d.Id(), shouldEnableRotation, err.Error())
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{fmt.Sprintf("%t", !shouldEnableRotation)},
		Target:                    []string{fmt.Sprintf("%t", shouldEnableRotation)},
		Timeout:                   5 * time.Minute,
		MinTimeout:                1 * time.Second,
		ContinuousTargetOccurence: 5,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s rotation status is %t",
				d.Id(), shouldEnableRotation)
			resp, err := conn.GetKeyRotationStatus(&kms.GetKeyRotationStatusInput{
				KeyId: aws.String(d.Id()),
			})
			if err != nil {
				return resp, "FAILED", err
			}
			status := fmt.Sprintf("%t", *resp.KeyRotationEnabled)
			log.Printf("[DEBUG] KMS key %s rotation status received: %s, retrying", d.Id(), status)

			return resp, status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed setting KMS key rotation status to %t: %s", shouldEnableRotation, err)
	}

	return nil
}

func resourceAwsKmsKeyDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn
	keyId := d.Get("key_id").(string)

	req := &kms.ScheduleKeyDeletionInput{
		KeyId: aws.String(keyId),
	}
	if v, exists := d.GetOk("deletion_window_in_days"); exists {
		req.PendingWindowInDays = aws.Int64(int64(v.(int)))
	}
	_, err := conn.ScheduleKeyDeletion(req)
	if err != nil {
		return err
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{"Enabled", "Disabled"},
		Target:                    []string{"PendingDeletion"},
		Timeout:                   20 * time.Minute,
		MinTimeout:                2 * time.Second,
		ContinuousTargetOccurence: 10,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s state is PendingDeletion", keyId)
			resp, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(keyId),
			})
			if err != nil {
				return resp, "Failed", err
			}

			metadata := *resp.KeyMetadata
			log.Printf("[DEBUG] KMS key %s state is %s, retrying", keyId, *metadata.KeyState)

			return resp, *metadata.KeyState, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed deactivating KMS key %s: %s", keyId, err)
	}

	log.Printf("[DEBUG] KMS Key %s deactivated.", keyId)
	d.SetId("")
	return nil
}
<commit_msg>Handle JSON parsing error in the ReadFunc for the policy document.<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsKmsKey() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsKmsKeyCreate,
		Read:   resourceAwsKmsKeyRead,
		Update: resourceAwsKmsKeyUpdate,
		Delete: resourceAwsKmsKeyDelete,

		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"key_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"key_usage": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, es []error) {
					value := v.(string)
					if !(value == "ENCRYPT_DECRYPT" || value == "") {
						es = append(es, fmt.Errorf(
							"%q must be ENCRYPT_DECRYPT or not specified", k))
					}
					return
				},
			},
			"policy": &schema.Schema{
				Type:             schema.TypeString,
				Optional:         true,
				Computed:         true,
				ValidateFunc:     validateJsonString,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},
			"is_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"enable_key_rotation": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},
			"deletion_window_in_days": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, es []error) {
					value := v.(int)
					if value > 30 || value < 7 {
						es = append(es, fmt.Errorf(
							"%q must be between 7 and 30 days inclusive", k))
					}
					return
				},
			},
		},
	}
}

func resourceAwsKmsKeyCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn

	// Allow aws to chose default values if we don't pass them
	var req kms.CreateKeyInput
	if v, exists := d.GetOk("description"); exists {
		req.Description = aws.String(v.(string))
	}
	if v, exists := d.GetOk("key_usage"); exists {
		req.KeyUsage = aws.String(v.(string))
	}
	if v, exists := d.GetOk("policy"); exists {
		req.Policy = aws.String(v.(string))
	}

	resp, err := conn.CreateKey(&req)
	if err != nil {
		return err
	}

	d.SetId(*resp.KeyMetadata.KeyId)
	d.Set("key_id", resp.KeyMetadata.KeyId)

	return _resourceAwsKmsKeyUpdate(d, meta, true)
}

func resourceAwsKmsKeyRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn

	req := &kms.DescribeKeyInput{
		KeyId: aws.String(d.Id()),
	}
	resp, err := conn.DescribeKey(req)
	if err != nil {
		return err
	}
	metadata := resp.KeyMetadata

	if *metadata.KeyState == "PendingDeletion" {
		log.Printf("[WARN] Removing KMS key %s because it's already gone", d.Id())
		d.SetId("")
		return nil
	}

	d.SetId(*metadata.KeyId)

	d.Set("arn", metadata.Arn)
	d.Set("key_id", metadata.KeyId)
	d.Set("description", metadata.Description)
	d.Set("key_usage", metadata.KeyUsage)
	d.Set("is_enabled", metadata.Enabled)

	p, err := conn.GetKeyPolicy(&kms.GetKeyPolicyInput{
		KeyId:      metadata.KeyId,
		PolicyName: aws.String("default"),
	})
	if err != nil {
		return err
	}

	policy, err := normalizeJsonString(*p.Policy)
	if err != nil {
		return errwrap.Wrapf("policy contains an invalid JSON: {{err}}", err)
	}
	d.Set("policy", policy)

	krs, err := conn.GetKeyRotationStatus(&kms.GetKeyRotationStatusInput{
		KeyId: metadata.KeyId,
	})
	if err != nil {
		return err
	}
	d.Set("enable_key_rotation", krs.KeyRotationEnabled)

	return nil
}

func resourceAwsKmsKeyUpdate(d *schema.ResourceData, meta interface{}) error {
	return _resourceAwsKmsKeyUpdate(d, meta, false)
}

// We expect new keys to be enabled already
// but there is no easy way to differentiate between Update()
// called from Create() and regular update, so we have this wrapper
func _resourceAwsKmsKeyUpdate(d *schema.ResourceData, meta interface{}, isFresh bool) error {
	conn := meta.(*AWSClient).kmsconn

	if d.HasChange("is_enabled") && d.Get("is_enabled").(bool) && !isFresh {
		// Enable before any attributes will be modified
		if err := updateKmsKeyStatus(conn, d.Id(), d.Get("is_enabled").(bool)); err != nil {
			return err
		}
	}

	if d.HasChange("enable_key_rotation") {
		if err := updateKmsKeyRotationStatus(conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("description") {
		if err := resourceAwsKmsKeyDescriptionUpdate(conn, d); err != nil {
			return err
		}
	}
	if d.HasChange("policy") {
		if err := resourceAwsKmsKeyPolicyUpdate(conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("is_enabled") && !d.Get("is_enabled").(bool) {
		// Only disable when all attributes are modified
		// because we cannot modify disabled keys
		if err := updateKmsKeyStatus(conn, d.Id(), d.Get("is_enabled").(bool)); err != nil {
			return err
		}
	}

	return resourceAwsKmsKeyRead(d, meta)
}

func resourceAwsKmsKeyDescriptionUpdate(conn *kms.KMS, d *schema.ResourceData) error {
	description := d.Get("description").(string)
	keyId := d.Get("key_id").(string)

	log.Printf("[DEBUG] KMS key: %s, update description: %s", keyId, description)

	req := &kms.UpdateKeyDescriptionInput{
		Description: aws.String(description),
		KeyId:       aws.String(keyId),
	}
	_, err := conn.UpdateKeyDescription(req)
	return err
}

func resourceAwsKmsKeyPolicyUpdate(conn *kms.KMS, d *schema.ResourceData) error {
	policy, err := normalizeJsonString(d.Get("policy").(string))
	if err != nil {
		return errwrap.Wrapf("policy contains an invalid JSON: {{err}}", err)
	}
	keyId := d.Get("key_id").(string)

	log.Printf("[DEBUG] KMS key: %s, update policy: %s", keyId, policy)

	req := &kms.PutKeyPolicyInput{
		KeyId:      aws.String(keyId),
		Policy:     aws.String(policy),
		PolicyName: aws.String("default"),
	}
	_, err = conn.PutKeyPolicy(req)
	return err
}

func updateKmsKeyStatus(conn *kms.KMS, id string, shouldBeEnabled bool) error {
	var err error

	if shouldBeEnabled {
		log.Printf("[DEBUG] Enabling KMS key %q", id)
		_, err = conn.EnableKey(&kms.EnableKeyInput{
			KeyId: aws.String(id),
		})
	} else {
		log.Printf("[DEBUG] Disabling KMS key %q", id)
		_, err = conn.DisableKey(&kms.DisableKeyInput{
			KeyId: aws.String(id),
		})
	}

	if err != nil {
		return fmt.Errorf("Failed to set KMS key %q status to %t: %q",
			id, shouldBeEnabled, err.Error())
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{fmt.Sprintf("%t", !shouldBeEnabled)},
		Target:                    []string{fmt.Sprintf("%t", shouldBeEnabled)},
		Timeout:                   20 * time.Minute,
		MinTimeout:                2 * time.Second,
		ContinuousTargetOccurence: 10,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s enabled status is %t",
				id, shouldBeEnabled)
			resp, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(id),
			})
			if err != nil {
				return resp, "FAILED", err
			}
			status := fmt.Sprintf("%t", *resp.KeyMetadata.Enabled)
			log.Printf("[DEBUG] KMS key %s status received: %s, retrying", id, status)

			return resp, status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed setting KMS key status to %t: %s", shouldBeEnabled, err)
	}

	return nil
}

func updateKmsKeyRotationStatus(conn *kms.KMS, d *schema.ResourceData) error {
	var err error
	shouldEnableRotation := d.Get("enable_key_rotation").(bool)
	if shouldEnableRotation {
		log.Printf("[DEBUG] Enabling key rotation for KMS key %q", d.Id())
		_, err = conn.EnableKeyRotation(&kms.EnableKeyRotationInput{
			KeyId: aws.String(d.Id()),
		})
	} else {
		log.Printf("[DEBUG] Disabling key rotation for KMS key %q", d.Id())
		_, err = conn.DisableKeyRotation(&kms.DisableKeyRotationInput{
			KeyId: aws.String(d.Id()),
		})
	}

	if err != nil {
		return fmt.Errorf("Failed to set key rotation for %q to %t: %q",
			d.Id(), shouldEnableRotation, err.Error())
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{fmt.Sprintf("%t", !shouldEnableRotation)},
		Target:                    []string{fmt.Sprintf("%t", shouldEnableRotation)},
		Timeout:                   5 * time.Minute,
		MinTimeout:                1 * time.Second,
		ContinuousTargetOccurence: 5,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s rotation status is %t",
				d.Id(), shouldEnableRotation)
			resp, err := conn.GetKeyRotationStatus(&kms.GetKeyRotationStatusInput{
				KeyId: aws.String(d.Id()),
			})
			if err != nil {
				return resp, "FAILED", err
			}
			status := fmt.Sprintf("%t", *resp.KeyRotationEnabled)
			log.Printf("[DEBUG] KMS key %s rotation status received: %s, retrying", d.Id(), status)

			return resp, status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed setting KMS key rotation status to %t: %s", shouldEnableRotation, err)
	}

	return nil
}

func resourceAwsKmsKeyDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn
	keyId := d.Get("key_id").(string)

	req := &kms.ScheduleKeyDeletionInput{
		KeyId: aws.String(keyId),
	}
	if v, exists := d.GetOk("deletion_window_in_days"); exists {
		req.PendingWindowInDays = aws.Int64(int64(v.(int)))
	}
	_, err := conn.ScheduleKeyDeletion(req)
	if err != nil {
		return err
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{"Enabled", "Disabled"},
		Target:                    []string{"PendingDeletion"},
		Timeout:                   20 * time.Minute,
		MinTimeout:                2 * time.Second,
		ContinuousTargetOccurence: 10,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s state is PendingDeletion", keyId)
			resp, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(keyId),
			})
			if err != nil {
				return resp, "Failed", err
			}

			metadata := *resp.KeyMetadata
			log.Printf("[DEBUG] KMS key %s state is %s, retrying", keyId, *metadata.KeyState)

			return resp, *metadata.KeyState, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed deactivating KMS key %s: %s", keyId, err)
	}

	log.Printf("[DEBUG] KMS Key %s deactivated.", keyId)
	d.SetId("")
	return nil
}
<|endoftext|>


<commit_before>package clients

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/amalgam8/controller/resources"
)

type NGINX interface {
	UpdateHttpUpstreams(conf resources.NGINXJson) error
}

type nginx struct {
	httpClient *http.Client
	url        string
}

func NewNGINXClient(url string) NGINX {
	return &nginx{
		httpClient: &http.Client{},
		url:        url,
	}
}

func (n *nginx) UpdateHttpUpstreams(conf resources.NGINXJson) error {

	data, err := json.Marshal(&conf)
	if err != nil {
		logrus.WithFields(logrus.Fields{
			"err": err,
		}).Error("Could not marshal request body")
		return err
	}

	reader := bytes.NewReader(data)
	req, err := http.NewRequest("POST", n.url+"/a8-admin", reader)
	if err != nil {
		logrus.WithError(err).Error("Failed building request to NGINX server")
		return err
	}

	resp, err := n.httpClient.Do(req)
	if err != nil {
		logrus.WithError(err).Error("Failed to send request to NGINX server")
	}

	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		data, _ := ioutil.ReadAll(resp.Body)

		logrus.WithFields(logrus.Fields{
			"err":         err,
			"body":        string(data),
			"status_code": resp.StatusCode,
		}).Error("POST to NGINX server return failure")
		return err
	}

	return nil
}
<commit_msg>bug fix<commit_after>package clients

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/amalgam8/controller/resources"
)

type NGINX interface {
	UpdateHttpUpstreams(conf resources.NGINXJson) error
}

type nginx struct {
	httpClient *http.Client
	url        string
}

func NewNGINXClient(url string) NGINX {
	return &nginx{
		httpClient: &http.Client{},
		url:        url,
	}
}

func (n *nginx) UpdateHttpUpstreams(conf resources.NGINXJson) error {

	data, err := json.Marshal(&conf)
	if err != nil {
		logrus.WithFields(logrus.Fields{
			"err": err,
		}).Error("Could not marshal request body")
		return err
	}

	reader := bytes.NewReader(data)
	req, err := http.NewRequest("POST", n.url+"/a8-admin", reader)
	if err != nil {
		logrus.WithError(err).Error("Failed building request to NGINX server")
		return err
	}

	resp, err := n.httpClient.Do(req)
	if err != nil {
		logrus.WithError(err).Error("Failed to send request to NGINX server")
		return err
	}

	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		data, _ := ioutil.ReadAll(resp.Body)

		logrus.WithFields(logrus.Fields{
			"err":         err,
			"body":        string(data),
			"status_code": resp.StatusCode,
		}).Error("POST to NGINX server return failure")
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package prifi

/*
* This is the internal part of the API. As probably the prifi-service will
* not have an external API, this will not have any API-functions.
 */

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"strconv"
	"errors"

	"github.com/dedis/cothority/app/lib/config"
	"github.com/dedis/cothority/log"
	"github.com/dedis/cothority/network"
	"github.com/dedis/cothority/sda"
	"github.com/lbarman/prifi_dev/sda/protocols"
)

// ServiceName is the name to refer to the Template service from another
// package.
const ServiceName = "PrifiService"

var serviceID sda.ServiceID

// Register Service with SDA
func init() {
	sda.RegisterNewService(ServiceName, newService)
	serviceID = sda.ServiceFactory.ServiceID(ServiceName)
}

// This struct contains the state of the service
type Service struct {
	// We need to embed the ServiceProcessor, so that incoming messages
	// are correctly handled.
	*sda.ServiceProcessor
	group *config.Group
	Storage *Storage
	path    string
	role prifi.PriFiRole
}

// This structure will be saved, on the contrary of the 'Service'-structure
// which has per-service information stored
type Storage struct {
	TrusteeID string
}

// StartTrustee has to take a configuration and start the necessary
// protocols to enable the trustee-mode.
func (s *Service) StartTrustee(group *config.Group) error {
	log.Info("Service", s, "running in trustee mode")
	s.group = group
	s.role = prifi.Trustee

	return nil
}

// StartRelay has to take a configuration and start the necessary
// protocols to enable the relay-mode.
func (s *Service) StartRelay(group *config.Group) error {
	log.Info("Service", s, "running in relay mode")
	s.group = group
	s.role = prifi.Relay

	var wrapper *prifi.PriFiSDAWrapper
	ids, relayIdentity := mapIdentities(group)

	// Start the PriFi protocol on a flat tree with the relay as root
	tree := group.Roster.GenerateNaryTreeWithRoot(100, &relayIdentity)
	pi, err := s.CreateProtocolService(prifi.ProtocolName, tree)

	if err != nil {
		log.Fatal("Unable to start Prifi protocol:", err)
	}

	// Assert that pi has type PriFiSDAWrapper
	wrapper = pi.(*prifi.PriFiSDAWrapper)

	wrapper.SetConfig(&prifi.PriFiSDAWrapperConfig{
		Identities: ids,
		Role: prifi.Relay,
	})
	wrapper.Start()

	return nil
}

// StartClient has to take a configuration and start the necessary
// protocols to enable the client-mode.
func (s *Service) StartClient(group *config.Group) error {
	log.Info("Service", s, "running in client mode")
	s.group = group
	s.role = prifi.Client

	return nil
}

// NewProtocol is called on all nodes of a Tree (except the root, since it is
// the one starting the protocol) so it's the Service that will be called to
// generate the PI on all others node.
// If you use CreateProtocolSDA, this will not be called, as the SDA will
// instantiate the protocol on its own. If you need more control at the
// instantiation of the protocol, use CreateProtocolService, and you can
// give some extra-configuration to your protocol in here.
func (s *Service) NewProtocol(tn *sda.TreeNodeInstance, conf *sda.GenericConfig) (sda.ProtocolInstance, error) {
	log.Lvl5("Setting node configuration from service")

	ids, _ := mapIdentities(s.group)

	pi, err := prifi.NewPriFiSDAWrapperProtocol(tn)
	if err != nil {
		return nil, err
	}

	// Assert that pi has type PriFiSDAWrapper
	wrapper := pi.(*prifi.PriFiSDAWrapper)

	wrapper.SetConfig(&prifi.PriFiSDAWrapperConfig{
		Identities: ids,
		Role: s.role,
	})

	return wrapper, nil
}

// saves the actual identity
func (s *Service) save() {
	log.Lvl3("Saving service")
	b, err := network.MarshalRegisteredType(s.Storage)
	if err != nil {
		log.Error("Couldn't marshal service:", err)
	} else {
		err = ioutil.WriteFile(s.path+"/prifi.bin", b, 0660)
		if err != nil {
			log.Error("Couldn't save file:", err)
		}
	}
}

// Tries to load the configuration and updates if a configuration
// is found, else it returns an error.
func (s *Service) tryLoad() error {
	configFile := s.path + "/identity.bin"
	b, err := ioutil.ReadFile(configFile)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("Error while reading %s: %s", configFile, err)
	}
	if len(b) > 0 {
		_, msg, err := network.UnmarshalRegistered(b)
		if err != nil {
			return fmt.Errorf("Couldn't unmarshal: %s", err)
		}
		log.Lvl3("Successfully loaded")
		s.Storage = msg.(*Storage)
	}
	return nil
}

// newService receives the context and a path where it can write its
// configuration, if desired. As we don't know when the service will exit,
// we need to save the configuration on our own from time to time.
func newService(c *sda.Context, path string) sda.Service {
	log.Info("Calling newService")
	s := &Service{
		ServiceProcessor: sda.NewServiceProcessor(c),
		path:             path,
	}
	if err := s.tryLoad(); err != nil {
		log.Error(err)
	}
	return s
}

// parseDescription extracts a PriFiIdentity from a string
func parseDescription(description string) (*prifi.PriFiIdentity, error) {
	desc := strings.Split(description, " ")
	if len(desc) == 1 && desc[0] == "relay" {
		return &prifi.PriFiIdentity{
			Role: prifi.Relay,
			Id: 0,
		}, nil
	} else if len(desc) == 2 {
		id, err := strconv.Atoi(desc[1]); if err != nil {
			return nil, errors.New("Unable to parse id:")
		} else {
			pid := prifi.PriFiIdentity{
				Id: id,
			}
			if desc[0] == "client" {
				pid.Role = prifi.Client
			} else if  desc[0] == "trustee" {
				pid.Role = prifi.Trustee
			} else {
				return nil, errors.New("Invalid role.")
			}
			return &pid, nil
		}
	} else {
		return nil, errors.New("Invalid description.")
	}
}

// mapIdentities reads the group configuration to assign PriFi roles
// to server identities and return them with the server
// identity of the relay.
func mapIdentities(group *config.Group) (map[network.Address]prifi.PriFiIdentity, network.ServerIdentity) {
	m := make(map[network.Address]prifi.PriFiIdentity)
	var relay network.ServerIdentity

	// Read the description of the nodes in the config file to assign them PriFi roles.
	nodeList := group.Roster.List
	for i := 0; i < len(nodeList); i++ {
		si := nodeList[i]
		id, err := parseDescription(group.GetDescription(si))
		if err != nil {
			log.Info("Cannot parse node description, skipping:", err)
		} else {
			m[si.Address] = *id
			if id.Role == prifi.Relay {
				relay = *si
			}
		}
	}

	// Check that there is exactly one relay and at least one trustee and client
	t, c, r := 0, 0, 0

	for _, v := range m {
		switch v.Role {
		case prifi.Relay: r++
		case prifi.Client: c++
		case prifi.Trustee: t++
		}
	}

	if !(t > 0 && c > 0 && r == 1) {
		log.Fatal("Config file does not contain exactly one relay and at least one trustee and client.")
	}

	return m, relay
}
<commit_msg>Added functions to send connection messages to the relay<commit_after>package prifi

/*
* This is the internal part of the API. As probably the prifi-service will
* not have an external API, this will not have any API-functions.
 */

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"strconv"
	"errors"

	"github.com/dedis/cothority/app/lib/config"
	"github.com/dedis/cothority/log"
	"github.com/dedis/cothority/network"
	"github.com/dedis/cothority/sda"
	"github.com/lbarman/prifi_dev/sda/protocols"
)

// ServiceName is the name to refer to the Template service from another
// package.
const ServiceName = "PrifiService"

var serviceID sda.ServiceID

// Register Service with SDA
func init() {
	sda.RegisterNewService(ServiceName, newService)
	serviceID = sda.ServiceFactory.ServiceID(ServiceName)
	network.RegisterPacketType(ConnectionResponse{})
}

// This struct contains the state of the service
type Service struct {
	// We need to embed the ServiceProcessor, so that incoming messages
	// are correctly handled.
	*sda.ServiceProcessor
	group *config.Group
	Storage *Storage
	path    string
	role prifi.PriFiRole
}

// This structure will be saved, on the contrary of the 'Service'-structure
// which has per-service information stored
type Storage struct {
	TrusteeID string
}

type ConnectionRequest struct {}
type DisconnectionRequest struct {}
type ConnectionResponse struct {
	Status bool
}

// HandleConnection receives connection requests from other nodes.
// It decides when another PriFi protocol should be started.
func (s *Service) HandleConnection(from *network.ServerIdentity, req *ConnectionRequest) (network.Body, error) {
	log.Info(">>>> Received connection request ! <<<<<")
	// TODO: Check that we are relay, store who wants to start the protocol and start if there are enough participants
	return &ConnectionResponse{true}, nil
}

// HandleDisconnection receives disconnection requests.
// It must stop the current PriFi protocol.
func (s *Service) HandleDisconnection(from *network.ServerIdentity, req *DisconnectionRequest) (network.Body, error) {
	// TODO: This one will be a bit more complicated
	return &ConnectionResponse{true}, nil
}

// StartTrustee has to take a configuration and start the necessary
// protocols to enable the trustee-mode.
func (s *Service) StartTrustee(group *config.Group) error {
	log.Info("Service", s, "running in trustee mode")
	s.group = group
	s.role = prifi.Trustee

	// Inform the relay that we want to join the protocol
	err := sendConnectionRequest(group)
	if err != nil {
		log.Error("Connection failed:", err)
	}

	return nil
}

// StartRelay has to take a configuration and start the necessary
// protocols to enable the relay-mode.
func (s *Service) StartRelay(group *config.Group) error {
	log.Info("Service", s, "running in relay mode")
	s.group = group
	s.role = prifi.Relay
	/*
	var wrapper *prifi.PriFiSDAWrapper
	ids, relayIdentity := mapIdentities(group)

	// Start the PriFi protocol on a flat tree with the relay as root
	tree := group.Roster.GenerateNaryTreeWithRoot(100, &relayIdentity)
	pi, err := s.CreateProtocolService(prifi.ProtocolName, tree)

	if err != nil {
		log.Fatal("Unable to start Prifi protocol:", err)
	}

	// Assert that pi has type PriFiSDAWrapper
	wrapper = pi.(*prifi.PriFiSDAWrapper)

	wrapper.SetConfig(&prifi.PriFiSDAWrapperConfig{
		Identities: ids,
		Role: prifi.Relay,
	})
	wrapper.Start()
	*/
	return nil
}

// StartClient has to take a configuration and start the necessary
// protocols to enable the client-mode.
func (s *Service) StartClient(group *config.Group) error {
	log.Info("Service", s, "running in client mode")
	s.group = group
	s.role = prifi.Client

	// Inform the relay that we want to join the protocol
	err := sendConnectionRequest(group)
	if err != nil {
		log.Error("Connection failed:", err)
	}

	return nil
}

// NewProtocol is called on all nodes of a Tree (except the root, since it is
// the one starting the protocol) so it's the Service that will be called to
// generate the PI on all others node.
// If you use CreateProtocolSDA, this will not be called, as the SDA will
// instantiate the protocol on its own. If you need more control at the
// instantiation of the protocol, use CreateProtocolService, and you can
// give some extra-configuration to your protocol in here.
func (s *Service) NewProtocol(tn *sda.TreeNodeInstance, conf *sda.GenericConfig) (sda.ProtocolInstance, error) {
	log.Lvl5("Setting node configuration from service")

	ids, _ := mapIdentities(s.group)

	pi, err := prifi.NewPriFiSDAWrapperProtocol(tn)
	if err != nil {
		return nil, err
	}

	// Assert that pi has type PriFiSDAWrapper
	wrapper := pi.(*prifi.PriFiSDAWrapper)

	wrapper.SetConfig(&prifi.PriFiSDAWrapperConfig{
		Identities: ids,
		Role: s.role,
	})

	return wrapper, nil
}

// saves the actual identity
func (s *Service) save() {
	log.Lvl3("Saving service")
	b, err := network.MarshalRegisteredType(s.Storage)
	if err != nil {
		log.Error("Couldn't marshal service:", err)
	} else {
		err = ioutil.WriteFile(s.path+"/prifi.bin", b, 0660)
		if err != nil {
			log.Error("Couldn't save file:", err)
		}
	}
}

// Tries to load the configuration and updates if a configuration
// is found, else it returns an error.
func (s *Service) tryLoad() error {
	configFile := s.path + "/identity.bin"
	b, err := ioutil.ReadFile(configFile)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("Error while reading %s: %s", configFile, err)
	}
	if len(b) > 0 {
		_, msg, err := network.UnmarshalRegistered(b)
		if err != nil {
			return fmt.Errorf("Couldn't unmarshal: %s", err)
		}
		log.Lvl3("Successfully loaded")
		s.Storage = msg.(*Storage)
	}
	return nil
}

// newService receives the context and a path where it can write its
// configuration, if desired. As we don't know when the service will exit,
// we need to save the configuration on our own from time to time.
func newService(c *sda.Context, path string) sda.Service {
	log.Lvl4("Calling newService")
	s := &Service{
		ServiceProcessor: sda.NewServiceProcessor(c),
		path:             path,
	}
	if err := s.tryLoad(); err != nil {
		log.Error(err)
	}
	if err:= s.RegisterMessages(s.HandleConnection, s.HandleDisconnection); err != nil {
		log.Fatal("Could not register handlers:", err)
	}
	return s
}

// parseDescription extracts a PriFiIdentity from a string
func parseDescription(description string) (*prifi.PriFiIdentity, error) {
	desc := strings.Split(description, " ")
	if len(desc) == 1 && desc[0] == "relay" {
		return &prifi.PriFiIdentity{
			Role: prifi.Relay,
			Id: 0,
		}, nil
	} else if len(desc) == 2 {
		id, err := strconv.Atoi(desc[1]); if err != nil {
			return nil, errors.New("Unable to parse id:")
		} else {
			pid := prifi.PriFiIdentity{
				Id: id,
			}
			if desc[0] == "client" {
				pid.Role = prifi.Client
			} else if  desc[0] == "trustee" {
				pid.Role = prifi.Trustee
			} else {
				return nil, errors.New("Invalid role.")
			}
			return &pid, nil
		}
	} else {
		return nil, errors.New("Invalid description.")
	}
}

// mapIdentities reads the group configuration to assign PriFi roles
// to server identities and return them with the server
// identity of the relay.
func mapIdentities(group *config.Group) (map[network.Address]prifi.PriFiIdentity, network.ServerIdentity) {
	m := make(map[network.Address]prifi.PriFiIdentity)
	var relay network.ServerIdentity

	// Read the description of the nodes in the config file to assign them PriFi roles.
	nodeList := group.Roster.List
	for i := 0; i < len(nodeList); i++ {
		si := nodeList[i]
		id, err := parseDescription(group.GetDescription(si))
		if err != nil {
			log.Info("Cannot parse node description, skipping:", err)
		} else {
			m[si.Address] = *id
			if id.Role == prifi.Relay {
				relay = *si
			}
		}
	}

	// Check that there is exactly one relay and at least one trustee and client
	t, c, r := 0, 0, 0

	for _, v := range m {
		switch v.Role {
		case prifi.Relay: r++
		case prifi.Client: c++
		case prifi.Trustee: t++
		}
	}

	if !(t > 0 && c > 0 && r == 1) {
		log.Fatal("Config file does not contain exactly one relay and at least one trustee and client.")
	}

	return m, relay
}

func sendConnectionRequest(group *config.Group) error {
	client := sda.NewClient(ServiceName)
	_, relayIdentity := mapIdentities(group)
	reply, err := client.Send(&relayIdentity, &ConnectionRequest{})
	if e := network.ErrMsg(reply, err); e != nil {
		return e
	} else {
		res := reply.Msg.(ConnectionResponse)

		if res.Status {
			return nil
		} else {
			return errors.New("Connection request refused by the relay.")
		}
	}
}
<|endoftext|>


<commit_before>package search

import (
	"errors"
	"math"
	"sort"

	gr "github.com/gonum/graph"
)

// Finds all shortest paths between start and goal
type AllPathFunc func(start, goal gr.Node) (path [][]gr.Node, cost float64, err error)

// Finds one path between start and goal, which it finds is arbitrary
type SinglePathFunc func(start, goal gr.Node) (path []gr.Node, cost float64, err error)

// This function returns two functions: one that will generate all shortest paths between two nodes with ids i and j, and one that will generate just one path.
//
// This algorithm requires the CrunchGraph interface which means it only works on nodes with dense ids since it uses an adjacency matrix.
//
// This algorithm isn't blazingly fast, but is relatively fast for the domain. It runs at O((number of vertices)^3), and successfully computes
// the cost between all pairs of vertices.
//
// Generating a single path should be pretty cheap after FW is done running. The AllPathFunc is likely to be considerably more expensive,
// simply because it has to effectively generate all combinations of known valid paths at each recursive step of the algorithm.
func FloydWarshall(graph gr.CrunchGraph, cost func(gr.Node, gr.Node) float64) (AllPathFunc, SinglePathFunc) {
	graph.Crunch()
	_, _, _, _, _, _, cost, _ = setupFuncs(graph, cost, nil)

	nodes := denseNodeSorter(graph.NodeList())
	sort.Sort(nodes)
	numNodes := len(nodes)

	dist := make([]float64, numNodes*numNodes)
	next := make([][]int, numNodes*numNodes)
	for i := 0; i < numNodes; i++ {
		for j := 0; j < numNodes; j++ {
			if j != i {
				dist[i+j*numNodes] = math.Inf(1)
			}
		}
	}

	edges := graph.EdgeList()
	for _, edge := range edges {
		u := edge.Head().ID()
		v := edge.Tail().ID()

		dist[u+v*numNodes] = cost(edge.Head(), edge.Tail())
	}

	for k := 0; k < numNodes; k++ {
		for i := 0; i < numNodes; i++ {
			for j := 0; j < numNodes; j++ {
				if dist[i+j*numNodes] > dist[i+k*numNodes]+dist[k+j*numNodes] {
					dist[i+j*numNodes] = dist[i+k*numNodes] + dist[k+j*numNodes]

					// Avoid generating too much garbage by reusing the memory in the list if we've allocated one already
					if next[i+j*numNodes] == nil {
						next[i+j*numNodes] = []int{k}
					} else {
						next[i+j*numNodes] = next[i+j*numNodes][:1]
						next[i+j*numNodes][0] = k
					}
					// If the cost between the nodes happens to be the same cost as what we know, add the approriate
					// intermediary to the list
				} else if math.Abs(dist[i+k*numNodes]+dist[k+j*numNodes]-dist[i+j*numNodes]) < 0.00001 {
					next[i+j*numNodes] = append(next[i+j*numNodes], k)
				}
			}
		}
	}

	return genAllPathsFunc(dist, next, nodes), genSinglePathFunc(dist, next, nodes)
}

func genAllPathsFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([][]gr.Node, float64, error) {
	numNodes := len(nodes)

	// A recursive function to reconstruct all possible paths.
	// It's not fast, but it's about as fast as can be reasonably expected
	var allPathFinder func(i, j int) ([][]gr.Node, error)
	allPathFinder = func(i, j int) ([][]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}
		intermediates := next[i+j*numNodes]
		if intermediates == nil {
			return [][]gr.Node{}, nil
		}

		toReturn := make([][]gr.Node, 0, len(intermediates))

		// This step is a tad convoluted: we have some list of intermediates.
		// We can think of each intermediate as a path junction
		//
		// At this junction, we can find all the shortest paths back to i,
		// and all the shortest paths down to j. Since this is a junction,
		// any predecessor path that runs through this intermediate may
		// freely choose any successor path to get to j. They'll all be
		// of equivalent length.
		//
		// Thus, for each intermediate, we run through and join each predecessor
		// path with each successor path via its junction.
		for _, intermediate := range intermediates {

			// Find predecessors
			preds, err := allPathFinder(i, intermediate)
			if err != nil {
				return nil, err
			}

			// Join each predecessor with its junction
			for a := range preds {
				preds[a] = append(preds[a], nodes[intermediate])
			}

			// Find successors
			succs, err := allPathFinder(intermediate, j)
			if err != nil {
				return nil, err
			}

			// Join each successor with its predecessor at the junction.
			// (the copying stuff is because slices are reference types)
			for a := range succs {
				for b := range preds {
					path := make([]gr.Node, len(succs[a]), len(succs[a])+len(preds[b]))
					copy(path, succs[a])
					path = append(path, preds[b]...)
					toReturn = append(toReturn, path)
				}
			}

		}

		return toReturn, nil
	}

	return func(start, goal gr.Node) ([][]gr.Node, float64, error) {
		paths, err := allPathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		for i := range paths {
			// Prepend start and postpend goal. pathFinder only does the intermediate steps
			paths[i] = append(paths[i], nil)
			copy(paths[i][1:], paths[i][:len(paths[i])-1])
			paths[i][0] = start
			paths[i] = append(paths[i], goal)
		}

		return paths, dist[start.ID()+goal.ID()*numNodes], nil
	}
}

func genSinglePathFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([]gr.Node, float64, error) {
	numNodes := len(nodes)

	var singlePathFinder func(i, j int) ([]gr.Node, error)
	singlePathFinder = func(i, j int) ([]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}

		intermediates := next[i+j*numNodes]
		if intermediates == nil {
			return []gr.Node{}, nil
		}

		intermediate := intermediates[0]
		path, err := singlePathFinder(i, intermediate)
		if err != nil {
			return nil, err
		}

		path = append(path, nodes[intermediate])
		p, err := singlePathFinder(intermediate, j)
		if err != nil {
			return nil, err
		}
		path = append(path, p...)

		return path, nil
	}

	return func(start, goal gr.Node) ([]gr.Node, float64, error) {
		path, err := singlePathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		path = append(path, nil)
		copy(path[1:], path[:len(path)-1])
		path[0] = start
		path = append(path, goal)

		return path, dist[start.ID()+goal.ID()*numNodes], nil
	}
}
<commit_msg>Added a note<commit_after>package search

import (
	"errors"
	"math"
	"sort"

	gr "github.com/gonum/graph"
)

// Finds all shortest paths between start and goal
type AllPathFunc func(start, goal gr.Node) (path [][]gr.Node, cost float64, err error)

// Finds one path between start and goal, which it finds is arbitrary
type SinglePathFunc func(start, goal gr.Node) (path []gr.Node, cost float64, err error)

// This function returns two functions that will generate all shortest paths between two nodes with ids i and j.
//
// This algorithm requires the CrunchGraph interface which means it only works on nodes with dense ids since it uses an adjacency matrix.
//
// This algorithm isn't blazingly fast, but is relatively fast for the domain. It runs at O((number of vertices)^3), and successfully computes
// the cost between all pairs of vertices. Using just a little extra memory, we can remember all shortest paths
func FloydWarshall(graph gr.CrunchGraph, cost func(gr.Node, gr.Node) float64) (AllPathFunc, SinglePathFunc) {
	graph.Crunch()
	_, _, _, _, _, _, cost, _ = setupFuncs(graph, cost, nil)

	nodes := denseNodeSorter(graph.NodeList())
	sort.Sort(nodes)
	numNodes := len(nodes)

	dist := make([]float64, numNodes*numNodes)
	next := make([][]int, numNodes*numNodes)
	for i := 0; i < numNodes; i++ {
		for j := 0; j < numNodes; j++ {
			if j != i {
				dist[i+j*numNodes] = math.Inf(1)
			}
		}
	}

	edges := graph.EdgeList()
	for _, edge := range edges {
		u := edge.Head().ID()
		v := edge.Tail().ID()

		dist[u+v*numNodes] = cost(edge.Head(), edge.Tail())
	}

	for k := 0; k < numNodes; k++ {
		for i := 0; i < numNodes; i++ {
			for j := 0; j < numNodes; j++ {
				if dist[i+j*numNodes] > dist[i+k*numNodes]+dist[k+j*numNodes] {
					dist[i+j*numNodes] = dist[i+k*numNodes] + dist[k+j*numNodes]

					// Avoid generating too much garbage by reusing the memory in the list if we've allocated one already
					if next[i+j*numNodes] == nil {
						next[i+j*numNodes] = []int{k}
					} else {
						next[i+j*numNodes] = next[i+j*numNodes][:1]
						next[i+j*numNodes][0] = k
					}
					// If the cost between the nodes happens to be the same cost as what we know, add the approriate
					// intermediary to the list
				} else if math.Abs(dist[i+k*numNodes]+dist[k+j*numNodes]-dist[i+j*numNodes]) < 0.00001 {
					next[i+j*numNodes] = append(next[i+j*numNodes], k)
				}
			}
		}
	}

	return genAllPathsFunc(dist, next, nodes), genSinglePathFunc(dist, next, nodes)
}

func genAllPathsFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([][]gr.Node, float64, error) {
	numNodes := len(nodes)

	// A recursive function to reconstruct all possible paths.
	// It's not fast, but it's about as fast as can be reasonably expected
	var allPathFinder func(i, j int) ([][]gr.Node, error)
	allPathFinder = func(i, j int) ([][]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}
		intermediates := next[i+j*numNodes]
		if intermediates == nil {
			return [][]gr.Node{}, nil
		}

		toReturn := make([][]gr.Node, 0, len(intermediates))

		// This step is a tad convoluted: we have some list of intermediates.
		// We can think of each intermediate as a path junction
		//
		// At this junction, we can find all the shortest paths back to i,
		// and all the shortest paths down to j. Since this is a junction,
		// any predecessor path that runs through this intermediate may
		// freely choose any successor path to get to j. They'll all be
		// of equivalent length.
		//
		// Thus, for each intermediate, we run through and join each predecessor
		// path with each successor path via its junction.
		for _, intermediate := range intermediates {

			// Find predecessors
			preds, err := allPathFinder(i, intermediate)
			if err != nil {
				return nil, err
			}

			// Join each predecessor with its junction
			for a := range preds {
				preds[a] = append(preds[a], nodes[intermediate])
			}

			// Find successors
			succs, err := allPathFinder(intermediate, j)
			if err != nil {
				return nil, err
			}

			// Join each successor with its predecessor at the junction.
			// (the copying stuff is because slices are reference types)
			for a := range succs {
				for b := range preds {
					path := make([]gr.Node, len(succs[a]), len(succs[a])+len(preds[b]))
					copy(path, succs[a])
					path = append(path, preds[b]...)
					toReturn = append(toReturn, path)
				}
			}

		}

		return toReturn, nil
	}

	return func(start, goal gr.Node) ([][]gr.Node, float64, error) {
		paths, err := allPathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		for i := range paths {
			// Prepend start and postpend goal. pathFinder only does the intermediate steps
			paths[i] = append(paths[i], nil)
			copy(paths[i][1:], paths[i][:len(paths[i])-1])
			paths[i][0] = start
			paths[i] = append(paths[i], goal)
		}

		return paths, dist[start.ID()+goal.ID()*numNodes], nil
	}
}

func genSinglePathFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([]gr.Node, float64, error) {
	numNodes := len(nodes)

	var singlePathFinder func(i, j int) ([]gr.Node, error)
	singlePathFinder = func(i, j int) ([]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}

		intermediates := next[i+j*numNodes]
		if intermediates == nil {
			return []gr.Node{}, nil
		}

		intermediate := intermediates[0]
		path, err := singlePathFinder(i, intermediate)
		if err != nil {
			return nil, err
		}

		path = append(path, nodes[intermediate])
		p, err := singlePathFinder(intermediate, j)
		if err != nil {
			return nil, err
		}
		path = append(path, p...)

		return path, nil
	}

	return func(start, goal gr.Node) ([]gr.Node, float64, error) {
		path, err := singlePathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		path = append(path, nil)
		copy(path[1:], path[:len(path)-1])
		path[0] = start
		path = append(path, goal)

		return path, dist[start.ID()+goal.ID()*numNodes], nil
	}
}
<|endoftext|>


<commit_before>package search

import (
	"errors"
	"math"
	"sort"

	gr "github.com/gonum/graph"
)

// Finds all shortest paths between start and goal
type AllPathFunc func(start, goal gr.Node) (path [][]gr.Node, cost float64, err error)

// Finds one path between start and goal, which it finds is arbitrary
type SinglePathFunc func(start, goal gr.Node) (path []gr.Node, cost float64, err error)

// This function returns two functions: one that will generate all shortest paths between two nodes with ids i and j, and one that will generate just one path.
//
// This algorithm requires the CrunchGraph interface which means it only works on nodes with dense ids since it uses an adjacency matrix.
//
// This algorithm isn't blazingly fast, but is relatively fast for the domain. It runs at O((number of vertices)^3), and successfully computes
// the cost between all pairs of vertices.
//
// Generating a single path should be pretty cheap after FW is done running. The AllPathFunc is likely to be considerably more expensive,
// simply because it has to effectively generate all combinations of known valid paths at each recursive step of the algorithm.
func FloydWarshall(graph gr.CrunchGraph, cost func(gr.Node, gr.Node) float64) (AllPathFunc, SinglePathFunc) {
	graph.Crunch()
	_, _, _, _, _, _, cost, _ = setupFuncs(graph, cost, nil)

	nodes := denseNodeSorter(graph.NodeList())
	sort.Sort(nodes)
	numNodes := len(nodes)

	dist := make([]float64, numNodes*numNodes)
	next := make([][]int, numNodes*numNodes)
	for i := 0; i < numNodes; i++ {
		for j := 0; j < numNodes; j++ {
			if j != i {
				dist[i+j*numNodes] = math.Inf(1)
			}
		}
	}

	edges := graph.EdgeList()
	for _, edge := range edges {
		u := edge.Head().ID()
		v := edge.Tail().ID()

		dist[u+v*numNodes] = cost(edge.Head(), edge.Tail())
	}

	for k := 0; k < numNodes; k++ {
		for i := 0; i < numNodes; i++ {
			for j := 0; j < numNodes; j++ {
				if dist[i+j*numNodes] > dist[i+k*numNodes]+dist[k+j*numNodes] {
					dist[i+j*numNodes] = dist[i+k*numNodes] + dist[k+j*numNodes]

					// Avoid generating too much garbage by reusing the memory in the list if we've allocated one already
					if next[i+j*numNodes] == nil {
						next[i+j*numNodes] = []int{k}
					} else {
						next[i+j*numNodes] = next[i+j*numNodes][:1]
						next[i+j*numNodes][0] = k
					}
					// If the cost between the nodes happens to be the same cost as what we know, add the approriate
					// intermediary to the list
				} else if math.Abs(dist[i+k*numNodes]+dist[k+j*numNodes]-dist[i+j*numNodes]) < 0.00001 {
					next[i+j*numNodes] = append(next[i+j*numNodes], k)
				}
			}
		}
	}

	return genAllPathsFunc(dist, next, nodes), genSinglePathFunc(dist, next, nodes)
}

func genAllPathsFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([][]gr.Node, float64, error) {
	numNodes := len(nodes)

	// A recursive function to reconstruct all possible paths.
	// It's not fast, but it's about as fast as can be reasonably expected
	var allPathFinder func(i, j int) ([][]gr.Node, error)
	allPathFinder = func(i, j int) ([][]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}
		intermediates := next[i+j*numNodes]
		if intermediates == nil {
			return [][]gr.Node{}, nil
		}

		toReturn := make([][]gr.Node, 0, len(intermediates))

		// This step is a tad convoluted: we have some list of intermediates.
		// We can think of each intermediate as a path junction
		//
		// At this junction, we can find all the shortest paths back to i,
		// and all the shortest paths down to j. Since this is a junction,
		// any predecessor path that runs through this intermediate may
		// freely choose any successor path to get to j. They'll all be
		// of equivalent length.
		//
		// Thus, for each intermediate, we run through and join each predecessor
		// path with each successor path via its junction.
		for _, intermediate := range intermediates {

			// Find predecessors
			preds, err := allPathFinder(i, intermediate)
			if err != nil {
				return nil, err
			}

			// Join each predecessor with its junction
			for a := range preds {
				preds[a] = append(preds[a], nodes[intermediate])
			}

			// Find successors
			succs, err := allPathFinder(intermediate, j)
			if err != nil {
				return nil, err
			}

			// Join each successor with its predecessor at the junction.
			// (the copying stuff is because slices are reference types)
			for a := range succs {
				for b := range preds {
					path := make([]gr.Node, len(succs[a]), len(succs[a])+len(preds[b]))
					copy(path, succs[a])
					path = append(path, preds[b]...)
					toReturn = append(toReturn, path)
				}
			}

		}

		return toReturn, nil
	}

	return func(start, goal gr.Node) ([][]gr.Node, float64, error) {
		paths, err := allPathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		for i := range paths {
			// Prepend start and postpend goal. pathFinder only does the intermediate steps
			paths[i] = append(paths[i], nil)
			copy(paths[i][1:], paths[i][:len(paths[i])-1])
			paths[i][0] = start
			paths[i] = append(paths[i], goal)
		}

		return paths, dist[start.ID()+goal.ID()*numNodes], nil
	}
}

func genSinglePathFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([]gr.Node, float64, error) {
	numNodes := len(nodes)

	var singlePathFinder func(i, j int) ([]gr.Node, error)
	singlePathFinder = func(i, j int) ([]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}

		intermediates := next[i+j*numNodes]
		if intermediates == nil {
			return []gr.Node{}, nil
		}

		intermediate := intermediates[0]
		path, err := singlePathFinder(i, intermediate)
		if err != nil {
			return nil, err
		}

		path = append(path, nodes[intermediate])
		p, err := singlePathFinder(intermediate, j)
		if err != nil {
			return nil, err
		}
		path = append(path, p...)

		return path, nil
	}

	return func(start, goal gr.Node) ([]gr.Node, float64, error) {
		path, err := singlePathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		path = append(path, nil)
		copy(path[1:], path[:len(path)-1])
		path[0] = start
		path = append(path, goal)

		return path, dist[start.ID()+goal.ID()*numNodes], nil
	}
}
<commit_msg>Added a note<commit_after>package search

import (
	"errors"
	"math"
	"sort"

	gr "github.com/gonum/graph"
)

// Finds all shortest paths between start and goal
type AllPathFunc func(start, goal gr.Node) (path [][]gr.Node, cost float64, err error)

// Finds one path between start and goal, which it finds is arbitrary
type SinglePathFunc func(start, goal gr.Node) (path []gr.Node, cost float64, err error)

// This function returns two functions: one that will generate all shortest paths between two nodes with ids i and j, and one that will generate just one path.
//
// This algorithm requires the CrunchGraph interface which means it only works on nodes with dense ids since it uses an adjacency matrix.
//
// This algorithm isn't blazingly fast, but is relatively fast for the domain. It runs at O((number of vertices)^3), and successfully computes
// the cost between all pairs of vertices.
//
// Generating a single path should be pretty cheap after FW is done running. The AllPathFunc is likely to be considerably more expensive,
// simply because it has to effectively generate all combinations of known valid paths at each recursive step of the algorithm.
func FloydWarshall(graph gr.CrunchGraph, cost func(gr.Node, gr.Node) float64) (AllPathFunc, SinglePathFunc) {
	graph.Crunch()
	_, _, _, _, _, _, cost, _ = setupFuncs(graph, cost, nil)

	nodes := denseNodeSorter(graph.NodeList())
	sort.Sort(nodes)
	numNodes := len(nodes)

	dist := make([]float64, numNodes*numNodes)
	next := make([][]int, numNodes*numNodes)
	for i := 0; i < numNodes; i++ {
		for j := 0; j < numNodes; j++ {
			if j != i {
				dist[i+j*numNodes] = math.Inf(1)
			}
		}
	}

	edges := graph.EdgeList()
	for _, edge := range edges {
		u := edge.Head().ID()
		v := edge.Tail().ID()

		dist[u+v*numNodes] = cost(edge.Head(), edge.Tail())
	}

	for k := 0; k < numNodes; k++ {
		for i := 0; i < numNodes; i++ {
			for j := 0; j < numNodes; j++ {
				if dist[i+j*numNodes] > dist[i+k*numNodes]+dist[k+j*numNodes] {
					dist[i+j*numNodes] = dist[i+k*numNodes] + dist[k+j*numNodes]

					// Avoid generating too much garbage by reusing the memory in the list if we've allocated one already
					if next[i+j*numNodes] == nil {
						next[i+j*numNodes] = []int{k}
					} else {
						next[i+j*numNodes] = next[i+j*numNodes][:1]
						next[i+j*numNodes][0] = k
					}
					// If the cost between the nodes happens to be the same cost as what we know, add the approriate
					// intermediary to the list
					//
					// NOTE: This may be a straight else, awaiting tests.
				} else if math.Abs(dist[i+k*numNodes]+dist[k+j*numNodes]-dist[i+j*numNodes]) < 0.00001 {
					next[i+j*numNodes] = append(next[i+j*numNodes], k)
				}
			}
		}
	}

	return genAllPathsFunc(dist, next, nodes), genSinglePathFunc(dist, next, nodes)
}

func genAllPathsFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([][]gr.Node, float64, error) {
	numNodes := len(nodes)

	// A recursive function to reconstruct all possible paths.
	// It's not fast, but it's about as fast as can be reasonably expected
	var allPathFinder func(i, j int) ([][]gr.Node, error)
	allPathFinder = func(i, j int) ([][]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}
		intermediates := next[i+j*numNodes]
		if intermediates == nil {
			return [][]gr.Node{}, nil
		}

		toReturn := make([][]gr.Node, 0, len(intermediates))

		// This step is a tad convoluted: we have some list of intermediates.
		// We can think of each intermediate as a path junction
		//
		// At this junction, we can find all the shortest paths back to i,
		// and all the shortest paths down to j. Since this is a junction,
		// any predecessor path that runs through this intermediate may
		// freely choose any successor path to get to j. They'll all be
		// of equivalent length.
		//
		// Thus, for each intermediate, we run through and join each predecessor
		// path with each successor path via its junction.
		for _, intermediate := range intermediates {

			// Find predecessors
			preds, err := allPathFinder(i, intermediate)
			if err != nil {
				return nil, err
			}

			// Join each predecessor with its junction
			for a := range preds {
				preds[a] = append(preds[a], nodes[intermediate])
			}

			// Find successors
			succs, err := allPathFinder(intermediate, j)
			if err != nil {
				return nil, err
			}

			// Join each successor with its predecessor at the junction.
			// (the copying stuff is because slices are reference types)
			for a := range succs {
				for b := range preds {
					path := make([]gr.Node, len(succs[a]), len(succs[a])+len(preds[b]))
					copy(path, succs[a])
					path = append(path, preds[b]...)
					toReturn = append(toReturn, path)
				}
			}

		}

		return toReturn, nil
	}

	return func(start, goal gr.Node) ([][]gr.Node, float64, error) {
		paths, err := allPathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		for i := range paths {
			// Prepend start and postpend goal. pathFinder only does the intermediate steps
			paths[i] = append(paths[i], nil)
			copy(paths[i][1:], paths[i][:len(paths[i])-1])
			paths[i][0] = start
			paths[i] = append(paths[i], goal)
		}

		return paths, dist[start.ID()+goal.ID()*numNodes], nil
	}
}

func genSinglePathFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([]gr.Node, float64, error) {
	numNodes := len(nodes)

	var singlePathFinder func(i, j int) ([]gr.Node, error)
	singlePathFinder = func(i, j int) ([]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}

		intermediates := next[i+j*numNodes]
		if intermediates == nil {
			return []gr.Node{}, nil
		}

		intermediate := intermediates[0]
		path, err := singlePathFinder(i, intermediate)
		if err != nil {
			return nil, err
		}

		path = append(path, nodes[intermediate])
		p, err := singlePathFinder(intermediate, j)
		if err != nil {
			return nil, err
		}
		path = append(path, p...)

		return path, nil
	}

	return func(start, goal gr.Node) ([]gr.Node, float64, error) {
		path, err := singlePathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		path = append(path, nil)
		copy(path[1:], path[:len(path)-1])
		path[0] = start
		path = append(path, goal)

		return path, dist[start.ID()+goal.ID()*numNodes], nil
	}
}
<|endoftext|>


<commit_before>package server

import (
	"go.uber.org/zap"
	"sync"
)

// Parent function that will update all of the defined configuration files for a server
// automatically to ensure that they always use the specified values.
func (s *Server) UpdateConfigurationFiles() {
	wg := new(sync.WaitGroup)

	for _, v := range s.processConfiguration.ConfigurationFiles {
		wg.Add(1)

		go func(server *Server) {
			defer wg.Done()

			p, err := s.Filesystem.SafePath(v.FileName)
			if err != nil {
				zap.S().Errorw("failed to generate safe path for configuration file", zap.String("server", server.Uuid), zap.Error(err))

				return
			}

			if err := v.Parse(p); err != nil {
				zap.S().Errorw("failed to parse and update server configuration file", zap.String("server", server.Uuid), zap.Error(err))
			}
		}(s)
	}

	wg.Wait()
}<commit_msg>Fix improper goroutine<commit_after>package server

import (
	"github.com/pterodactyl/wings/parser"
	"go.uber.org/zap"
	"sync"
)

// Parent function that will update all of the defined configuration files for a server
// automatically to ensure that they always use the specified values.
func (s *Server) UpdateConfigurationFiles() {
	wg := new(sync.WaitGroup)

	for _, v := range s.processConfiguration.ConfigurationFiles {
		wg.Add(1)

		go func(f parser.ConfigurationFile, server *Server) {
			defer wg.Done()

			p, err := s.Filesystem.SafePath(f.FileName)
			if err != nil {
				zap.S().Errorw("failed to generate safe path for configuration file", zap.String("server", server.Uuid), zap.Error(err))

				return
			}

			if err := f.Parse(p); err != nil {
				zap.S().Errorw("failed to parse and update server configuration file", zap.String("server", server.Uuid), zap.Error(err))
			}
		}(v, s)
	}

	wg.Wait()
}<|endoftext|>


<commit_before>package server

import "strings"

var fatalError string = "fatal: invalid GroupedError"

type GroupedError struct {
	// The prefix string returned by `Error()`, followed by the grouped errors.
	Prefix string
	Errors []error
}

func (gErr *GroupedError) Error() string {
	if len(gErr.Errors) == 0 {
		return fatalError
	}
	var sb strings.Builder
	for _, err := range gErr.Errors {
		sb.WriteString("\n")
		sb.WriteString(err.Error())
	}
	return gErr.Prefix + sb.String()
}

func createGroupedError(prefix string, errors []error) error {
	if len(errors) == 0 {
		return nil
	}
	return &GroupedError{Prefix: prefix, Errors: errors}
}
<commit_msg>Add helper methods for dealing with GroupedErrors<commit_after>package server

import "strings"

var fatalError string = "fatal: invalid GroupedError"

type GroupedError struct {
	// The prefix string returned by `Error()`, followed by the grouped errors.
	Prefix string
	Errors []error
}

func (gErr *GroupedError) Error() string {
	if len(gErr.Errors) == 0 {
		return fatalError
	}
	var sb strings.Builder
	for _, err := range gErr.Errors {
		sb.WriteString("\n")
		sb.WriteString(err.Error())
	}
	return gErr.Prefix + sb.String()
}

func createGroupedError(prefix string, errors []error) error {
	if len(errors) == 0 {
		return nil
	}
	return &GroupedError{Prefix: prefix, Errors: errors}
}

func (gErr *GroupedError) containsSubstring(substr string) bool {
	for _, err := range gErr.Errors {
		if strings.Contains(err.Error(), substr) {
			return true
		}
	}
	return false
}

func (gErr *GroupedError) containsOnlySubstring(substr string) bool {
	if len(gErr.Errors) != 1 {
		return false
	}
	return gErr.containsSubstring(substr)
}
<|endoftext|>


<commit_before>package mcstore

import (
	"encoding/json"
	"net/http"
	"net/url"
	"strings"

	"github.com/materials-commons/config"
	"github.com/materials-commons/gohandy/ezhttp"
	"github.com/materials-commons/mcstore/pkg/app"
	"gnd.la/net/urlutil"
)

// MCUrl returns the current mcurl config entry.
func MCUrl() string {
	return config.GetString("mcurl")
}

// MCClient creates a new EzClient.
func MCClient() *ezhttp.EzClient {
	mcurl := MCUrl()
	if strings.HasPrefix(mcurl, "https") {
		return ezhttp.NewSSLClient()
	}
	return ezhttp.NewClient()
}

// Url create the url for accessing a service. It adds the mcurl to
// the path, and also adds the apikey argument.
func Url(path string) string {
	values := url.Values{}
	values.Add("apikey", config.GetString("apikey"))
	mcurl := urlutil.MustJoin(MCUrl(), path)
	mcurl = urlutil.AppendQuery(mcurl, values)
	return mcurl
}

// ToError tests the list of errors and the response to determine
// the type of error to return. It calls HTTPStatusToError to
// translate response status codes to an error.
func ToError(resp *http.Response, errs []error) error {
	if len(errs) != 0 {
		return app.ErrInvalid
	}
	return HTTPStatusToError(resp.StatusCode)
}

// HTTPStatusToError translates an http state to an
// application error.
func HTTPStatusToError(status int) error {
	switch {
	case status == http.StatusInternalServerError:
		return app.ErrInternal
	case status == http.StatusBadRequest:
		return app.ErrInvalid
	case status == http.StatusNotFound:
		return app.ErrNotFound
	case status == http.StatusForbidden:
		return app.ErrExists
	case status == http.StatusUnauthorized:
		return app.ErrNoAccess
	case status > 299:
		app.Log.Errorf("Unclassified error %d", status)
		return app.ErrUnclassified
	default:
		return nil
	}
}

// ToJSON unmarshalls a string that contains JSON.
func ToJSON(from string, to interface{}) error {
	err := json.Unmarshal([]byte(from), to)
	return err
}
<commit_msg>Fix up url handling since utilutil.Join wasn't doing the right thing.<commit_after>package mcstore

import (
	"encoding/json"
	"net/http"
	"net/url"
	"strings"

	"github.com/materials-commons/config"
	"github.com/materials-commons/gohandy/ezhttp"
	"github.com/materials-commons/mcstore/pkg/app"
	"gnd.la/net/urlutil"
)

// MCUrl returns the current mcurl config entry.
func MCUrl() string {
	return config.GetString("mcurl")
}

// MCClient creates a new EzClient.
func MCClient() *ezhttp.EzClient {
	mcurl := MCUrl()
	if strings.HasPrefix(mcurl, "https") {
		return ezhttp.NewSSLClient()
	}
	return ezhttp.NewClient()
}

// Url create the url for accessing a service. It adds the mcurl to
// the path, and also adds the apikey argument.
func Url(path string) string {
	values := url.Values{}
	values.Add("apikey", config.GetString("apikey"))
	mcurl := MCUrl()
	if strings.HasPrefix("/", path) {
		mcurl = MCUrl() + path
	} else {
		mcurl = MCUrl() + "/" + path
	}
	mcurl = urlutil.AppendQuery(mcurl, values)
	return mcurl
}

// ToError tests the list of errors and the response to determine
// the type of error to return. It calls HTTPStatusToError to
// translate response status codes to an error.
func ToError(resp *http.Response, errs []error) error {
	if len(errs) != 0 {
		return app.ErrInvalid
	}
	return HTTPStatusToError(resp.StatusCode)
}

// HTTPStatusToError translates an http state to an
// application error.
func HTTPStatusToError(status int) error {
	switch {
	case status == http.StatusInternalServerError:
		return app.ErrInternal
	case status == http.StatusBadRequest:
		return app.ErrInvalid
	case status == http.StatusNotFound:
		return app.ErrNotFound
	case status == http.StatusForbidden:
		return app.ErrExists
	case status == http.StatusUnauthorized:
		return app.ErrNoAccess
	case status > 299:
		app.Log.Errorf("Unclassified error %d", status)
		return app.ErrUnclassified
	default:
		return nil
	}
}

// ToJSON unmarshalls a string that contains JSON.
func ToJSON(from string, to interface{}) error {
	err := json.Unmarshal([]byte(from), to)
	return err
}
<|endoftext|>


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

package source

import (
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestFileInfo(t *testing.T) {
	assert := require.New(t)

	s := newTestSourceSpec()

	for _, this := range []struct {
		base     string
		filename string
		assert   func(f *FileInfo)
	}{
		{"/a/", filepath.FromSlash("/a/b/page.md"), func(f *FileInfo) {
			assert.Equal(filepath.FromSlash("/a/b/page.md"), f.Filename())
			assert.Equal(filepath.FromSlash("b/"), f.Dir())
			assert.Equal(filepath.FromSlash("b/page.md"), f.Path())

		}},
	} {
		f := s.NewFileInfo(this.base, this.filename, nil)
		this.assert(f)
	}

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

package source

import (
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestFileInfo(t *testing.T) {
	assert := require.New(t)

	s := newTestSourceSpec()

	for _, this := range []struct {
		base     string
		filename string
		assert   func(f *FileInfo)
	}{
		{filepath.FromSlash("/a/"), filepath.FromSlash("/a/b/page.md"), func(f *FileInfo) {
			assert.Equal(filepath.FromSlash("/a/b/page.md"), f.Filename())
			assert.Equal(filepath.FromSlash("b/"), f.Dir())
			assert.Equal(filepath.FromSlash("b/page.md"), f.Path())

		}},
	} {
		f := s.NewFileInfo(this.base, this.filename, nil)
		this.assert(f)
	}

}
<|endoftext|>


<commit_before>package database

import (
	"database/sql"
	"github.com/jmoiron/sqlx"
	"strings"
)

func makeProjectTable(db *sqlx.DB) error {
	_, err := db.Exec("CREATE TABLE IF NOT EXISTS projects (title TEXT NOT NULL PRIMARY KEY, screenshots TEXT NOT NULL, language TEXT NOT NULL, shortdesc TEXT NOT NULL, description TEXT NOT NULL)")

	return err
}

type Project struct {
	Title       string   // The title of the project
	Screenshots []string // URLs to screenshots of the project in action
	Language    string   // The language the project was written in
	ShortDesc   string   // A short version of the description
	Description string   // A description of the project
}

// Making a new project object
func MakeProject(title string, screenshots []string, language string, shortdesc string, description string) *Project {
	return &Project{
		Title:       title,
		Screenshots: screenshots,
		Language:    language,
		ShortDesc:   shortdesc,
		Description: description,
	}
}

// Making a list of projects from a sqlx.Rows object
func makeProjects(rows *sql.Rows) ([]*Project, error) {
	projects := make([]*Project, 0)
	var title string
	var screenshots string
	var language string
	var shortdesc string
	var description string

	for rows.Next() {
		err := rows.Scan(&title, &screenshots, &language, &shortdesc, &description)

		if err != nil {
			return nil, err
		}

		projects = append(projects, MakeProject(title, strings.Split(screenshots, ","), language, shortdesc, description))
	}

	return projects, nil
}

// Getting all of the projects
func GetProjects(db *sqlx.DB) ([]*Project, error) {
	rows, err := db.Query("SELECT * FROM projects")

	if err != nil {
		return nil, err
	}

	return makeProjects(rows)
}

// Getting a project by its title
func GetProjectByTitle(db *sqlx.DB, title string) (*Project, error) {
	stmt, err := db.Prepare("SELECT * FROM projects WHERE title = $1")

	if err != nil {
		return nil, err
	}

	row := stmt.QueryRow(title)

	var ntitle string
	var sscreenshots string
	var language string
	var shortdesc string
	var description string

	err = row.Scan(&ntitle, &sscreenshots, &language, &shortdesc, &description)

	var screenshots []string
	if sscreenshots == "" {
		screenshots = nil
	} else {
		screenshots = strings.Split(sscreenshots, ",")
	}

	return &Project{
		Title:       ntitle,
		Screenshots: screenshots,
		Language:    language,
		ShortDesc:   shortdesc,
		Description: description,
	}, nil
}

// Querying rows with a 'WHERE' statement
func queryProjectWithWhere(db *sqlx.DB, field string, value string) ([]*Project, error) {
	stmt, err := db.Prepare("SELECT * FROM projects WHERE $1 = $2")

	if err != nil {
		return nil, err
	}

	rows, err := stmt.Query(field, value)

	if err != nil {
		return nil, err
	}

	return makeProjects(rows)
}

// Querying a project by language
func QueryProjectByLanguage(db *sqlx.DB, language string) ([]*Project, error) {
	return queryProjectWithWhere(db, "language", language)
}
<commit_msg>Added functions for inserting posts and getting posts quickly.<commit_after>package database

import (
	"database/sql"
	"github.com/jmoiron/sqlx"
	"strings"
)

func makeProjectTable(db *sqlx.DB) error {
	_, err := db.Exec("CREATE TABLE IF NOT EXISTS projects (title TEXT NOT NULL PRIMARY KEY, screenshots TEXT NOT NULL, language TEXT NOT NULL, shortdesc TEXT NOT NULL, description TEXT NOT NULL)")

	return err
}

type Project struct {
	Title       string   // The title of the project
	Screenshots []string // URLs to screenshots of the project in action
	Language    string   // The language the project was written in
	ShortDesc   string   // A short version of the description
	Description string   // A description of the project
}

// Making a new project object
func MakeProject(title string, screenshots []string, language string, shortdesc string, description string) *Project {
	return &Project{
		Title:       title,
		Screenshots: screenshots,
		Language:    language,
		ShortDesc:   shortdesc,
		Description: description,
	}
}

// Making a list of projects from a sqlx.Rows object
func makeProjects(rows *sql.Rows) ([]*Project, error) {
	projects := make([]*Project, 0)
	var title string
	var screenshots string
	var language string
	var shortdesc string
	var description string

	for rows.Next() {
		err := rows.Scan(&title, &screenshots, &language, &shortdesc, &description)

		if err != nil {
			return nil, err
		}

		projects = append(projects, MakeProject(title, strings.Split(screenshots, ","), language, shortdesc, description))
	}

	return projects, nil
}

// Getting all of the projects
func GetProjects(db *sqlx.DB) ([]*Project, error) {
	rows, err := db.Query("SELECT * FROM projects")

	if err != nil {
		return nil, err
	}

	return makeProjects(rows)
}

// Quickly getting all of the projects
func QuickGetProjects() []*Project {
	db := QuickOpenDB()
	defer db.Close()

	projects, err := GetProjects(db)

	if err != nil {
		panic(err)
	}

	return projects
}

// Getting a project by its title
func GetProjectByTitle(db *sqlx.DB, title string) (*Project, error) {
	stmt, err := db.Prepare("SELECT * FROM projects WHERE title = $1")

	if err != nil {
		return nil, err
	}

	row := stmt.QueryRow(title)

	var ntitle string
	var sscreenshots string
	var language string
	var shortdesc string
	var description string

	err = row.Scan(&ntitle, &sscreenshots, &language, &shortdesc, &description)

	var screenshots []string
	if sscreenshots == "" {
		screenshots = nil
	} else {
		screenshots = strings.Split(sscreenshots, ",")
	}

	return &Project{
		Title:       ntitle,
		Screenshots: screenshots,
		Language:    language,
		ShortDesc:   shortdesc,
		Description: description,
	}, nil
}

// Quickly getting a project by its title
func QuickGetProjectByTitle(title string) *Project {
	db := QuickOpenDB()
	defer db.Close()

	project, err := GetProjectByTitle(db, title)

	if err != nil {
		panic(err)
	}

	return project
}

// Inserting a post
func InsertProject(db *sqlx.DB, project *Project) error {
	var joined string
	if len(project.Screenshots) == 0 {
		joined = ""
	} else {
		joined = project.Screenshots[0]
		for i := 1; i < len(project.Screenshots); i++ {
			joined += "," + project.Screenshots[i]
		}
	}

	exec := "INSERT INTO projects(title, screenshots, language, shortdesc, description) values($1, $2, $3, $4, $5)"

	_, err := db.Exec(exec, project.Title, joined, project.Language, project.ShortDesc, project.Description)
	return err
}

// Quickly inserting a post
func QuickInsertProject(project *Project) {
	db := QuickOpenDB()
	defer db.Close()

	if err := InsertProject(db, project); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package elasticthought

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/couchbaselabs/logg"
)

// Worker job that splits a dataset into training/test set
type DatasetSplitter struct {
	Configuration Configuration
	Dataset       Dataset
}

func (d DatasetSplitter) Run() {

	logg.LogTo("DATASET_SPLITTER", "Datasetsplitter.run()!.  Config: %+v Dataset: %+v", d.Configuration, d.Dataset)

	// Find the datafile object associated with dataset
	db := d.Configuration.DbConnection()
	datafile, err := d.Dataset.GetDatafile(db)
	if err != nil {
		errMsg := fmt.Errorf("Error looking up datafile: %v", err)
		logg.LogError(errMsg)
		return
	}

	// Open the url -- content type should be application/x-gzip and url should end with
	// .tar.gz
	tr1, tr2, err := d.openTwoTarGzStreams(datafile.Url)
	if err != nil {
		errMsg := fmt.Errorf("Error opening tar.gz streams: %v", err)
		logg.LogError(errMsg)
		return
	}

	cbfsReaderTesting, wTesting := io.Pipe()
	cbfsReaderTraining, wTraining := io.Pipe()

	go func() {
		bytes, err := ioutil.ReadAll(cbfsReaderTesting)
		logg.LogTo("DATASET_SPLITTER", "Read %d bytes from cbfsReaderTesting.  Err: %v", len(bytes), err)
	}()
	go func() {
		bytes, err := ioutil.ReadAll(cbfsReaderTraining)
		logg.LogTo("DATASET_SPLITTER", "Read %d bytes from cbfsReaderTraining.  Err: %v", len(bytes), err)
	}()

	tarWriterTesting := tar.NewWriter(wTesting)
	tarWriterTraining := tar.NewWriter(wTraining)

	err = d.transform(tr1, tr2, tarWriterTraining, tarWriterTesting)
	if err != nil {
		errMsg := fmt.Errorf("Error transforming tar stream: %v", err)
		logg.LogError(errMsg)
		return
	}

	// Read from the stream and open tar archive

	// Walk the directories and split the files

	// Write to training and test tar archive

	// Save both training and test tar archive to cbfs (wrapped in gzip stream)

}

// Opens to tar.gz streams to the same url.  The reason this is done twice is due to
// ugly hack, which is documented in the transform() method.
func (d DatasetSplitter) openTwoTarGzStreams(url string) (*tar.Reader, *tar.Reader, error) {

	resp1, err := http.Get(url)
	if err != nil {
		return nil, nil, err
	}
	gzipReader1, err := gzip.NewReader(resp1.Body)
	if err != nil {
		return nil, nil, err
	}
	tarReader1 := tar.NewReader(gzipReader1)

	resp2, err := http.Get(url)
	if err != nil {
		return nil, nil, err
	}
	gzipReader2, err := gzip.NewReader(resp2.Body)
	if err != nil {
		return nil, nil, err
	}
	tarReader2 := tar.NewReader(gzipReader2)

	return tarReader1, tarReader2, nil

}

// Read from source tar stream and write training and test to given tar writers
//
// TODO: fix ugly hack.  Since I'm trying to read from the source stream *twice*, which
// doesn't work, the workaround is to expect *two* source streams: source1 and source2.
// That way after source1 is read, source2 is ready for reading from the beginning
func (d DatasetSplitter) transform(source1, source2 *tar.Reader, train, test *tar.Writer) error {

	// build a map from the source
	sourceMap, err := d.createMap(source1)
	if err != nil {
		return err
	}
	logg.Log("sourceMap: %+v", sourceMap)

	// split the map into training and test
	trainMap, testMap, err := d.splitMap(sourceMap)
	if err != nil {
		return err
	}

	// iterate over the source
	logg.Log("iterate over source")
	for {
		hdr, err := source2.Next()
		logg.Log("hdr: %v", hdr)

		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return err
		}

		// distribute to writers based on training and test maps
		var twToAdd *tar.Writer

		// if strings.HasPrefix(hdr.Name, "foo") {
		if trainMap.hasPath(hdr.Name) {
			// add to training tar writer
			twToAdd = train
		} else if testMap.hasPath(hdr.Name) {
			// add to testing tar writer
			twToAdd = test
		} else {
			logg.LogPanic("File not present in either test/train: %v", hdr.Name)
		}

		// TODO: is there a more efficient way to do this?
		bytes, err := ioutil.ReadAll(source2)
		if err != nil {
			return err
		}

		logg.Log("file: %v numbytes: %v", hdr.Name, len(bytes))

		hdrToAdd := &tar.Header{
			Name: hdr.Name,
			Size: int64(len(bytes)),
		}
		if err := twToAdd.WriteHeader(hdrToAdd); err != nil {
			return err
		}
		if _, err := twToAdd.Write(bytes); err != nil {
			return err
		}

	}

	// close writers
	if err := train.Close(); err != nil {
		return err
	}
	if err := test.Close(); err != nil {
		return err
	}

	return nil
}

// Validate that the source tar stream conforms to expected specs
func (d DatasetSplitter) validate(source *tar.Reader) (bool, error) {

	// validation rules:
	// 1. has at least 2 files
	// 2. the depth of each file is 2 (folder/filename.xxx)

	numFiles := 0
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return false, err
		}
		numFiles += 1

		pathComponents := strings.Split(hdr.Name, "/")
		if len(pathComponents) != 2 {
			return false, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

	}

	if numFiles < 2 {
		return false, fmt.Errorf("Archive must contain at least 2 files")
	}

	return true, nil
}

// Create a map of folder -> []filename for all entries in the archive
func (d DatasetSplitter) createMap(source *tar.Reader) (filemap, error) {

	resultMap := filemap{}
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return nil, err
		}

		pathComponents := strings.Split(hdr.Name, "/")

		if len(pathComponents) != 2 {
			return nil, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

		directory := pathComponents[0]
		filename := pathComponents[1]

		resultMap.addFileToDirectory(directory, filename)

	}

	return resultMap, nil

}

// Split map into training and testing disjoint subsets based on values
// of DatasetSplitter's Dataset
func (d DatasetSplitter) splitMap(source filemap) (training filemap, testing filemap, err error) {

	training = filemap{}
	testing = filemap{}

	// iterate over source keys
	for directory, files := range source {

		numTraining := int(float64(len(files)) * d.Dataset.TrainingDataset.SplitPercentage)

		numTest := len(files) - int(numTraining)

		// split files into subsets based on ratios in dataset
		trainingFiles, testFiles, err := splitFiles(files, numTraining, numTest)

		if err != nil {
			return nil, nil, err
		}

		// add to respective maps
		training[directory] = trainingFiles
		testing[directory] = testFiles

	}

	return training, testing, nil

}

func splitFiles(files []string, numTraining, numTest int) (training []string, test []string, err error) {
	for i, file := range files {
		if i < numTraining {
			training = append(training, file)
		} else {
			test = append(test, file)
		}
	}
	return
}
<commit_msg>add piped reader/writer, seems to be working<commit_after>package elasticthought

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/couchbaselabs/logg"
)

// Worker job that splits a dataset into training/test set
type DatasetSplitter struct {
	Configuration Configuration
	Dataset       Dataset
}

func (d DatasetSplitter) Run() {

	logg.LogTo("DATASET_SPLITTER", "Datasetsplitter.run()!.  Config: %+v Dataset: %+v", d.Configuration, d.Dataset)

	// Find the datafile object associated with dataset
	db := d.Configuration.DbConnection()
	datafile, err := d.Dataset.GetDatafile(db)
	if err != nil {
		errMsg := fmt.Errorf("Error looking up datafile with id: %v.  Error: %v", d.Dataset.DatafileID, err)
		logg.LogError(errMsg)
		return
	}

	// Open the url -- content type should be application/x-gzip and url should end with
	// .tar.gz
	tr1, tr2, err := d.openTwoTarGzStreams(datafile.Url)
	if err != nil {
		errMsg := fmt.Errorf("Error opening tar.gz streams: %v", err)
		logg.LogError(errMsg)
		return
	}

	logg.LogTo("DATASET_SPLITTER", "Creating io pipe")

	cbfsReaderTesting, wTesting := io.Pipe()
	cbfsReaderTraining, wTraining := io.Pipe()

	tarWriterTesting := tar.NewWriter(wTesting)
	tarWriterTraining := tar.NewWriter(wTraining)

	go func() {
		for {
			bytes := make([]byte, 1024)
			logg.LogTo("DATASET_SPLITTER", "Going to read bytes from cbfsReaderTesting.")
			n, err := cbfsReaderTesting.Read(bytes)
			logg.LogTo("DATASET_SPLITTER", "Read %d bytes from cbfsReaderTesting.  Err: %v", n, err)
			if err != nil {
				break
			}
		}

	}()
	go func() {
		for {
			bytes := make([]byte, 1024)
			logg.LogTo("DATASET_SPLITTER", "Going to read bytes from cbfsReaderTraining.")
			n, err := cbfsReaderTraining.Read(bytes)
			logg.LogTo("DATASET_SPLITTER", "Read %d bytes from cbfsReaderTraining.  Err: %v", n, err)
			if err != nil {
				break
			}

		}

	}()

	logg.LogTo("DATASET_SPLITTER", "Calling transform")
	err = d.transform(tr1, tr2, tarWriterTraining, tarWriterTesting)
	if err != nil {
		errMsg := fmt.Errorf("Error transforming tar stream: %v", err)
		logg.LogError(errMsg)
		return
	}

	logg.LogTo("DATASET_SPLITTER", "Finished calling transform")

	cbfsReaderTesting.Close()
	cbfsReaderTraining.Close()

	// Read from the stream and open tar archive

	// Walk the directories and split the files

	// Write to training and test tar archive

	// Save both training and test tar archive to cbfs (wrapped in gzip stream)

}

// Opens to tar.gz streams to the same url.  The reason this is done twice is due to
// ugly hack, which is documented in the transform() method.
func (d DatasetSplitter) openTwoTarGzStreams(url string) (*tar.Reader, *tar.Reader, error) {

	resp1, err := http.Get(url)
	if err != nil {
		return nil, nil, err
	}
	gzipReader1, err := gzip.NewReader(resp1.Body)
	if err != nil {
		return nil, nil, err
	}
	tarReader1 := tar.NewReader(gzipReader1)

	resp2, err := http.Get(url)
	if err != nil {
		return nil, nil, err
	}
	gzipReader2, err := gzip.NewReader(resp2.Body)
	if err != nil {
		return nil, nil, err
	}
	tarReader2 := tar.NewReader(gzipReader2)

	return tarReader1, tarReader2, nil

}

// Read from source tar stream and write training and test to given tar writers
//
// TODO: fix ugly hack.  Since I'm trying to read from the source stream *twice*, which
// doesn't work, the workaround is to expect *two* source streams: source1 and source2.
// That way after source1 is read, source2 is ready for reading from the beginning
func (d DatasetSplitter) transform(source1, source2 *tar.Reader, train, test *tar.Writer) error {

	// build a map from the source
	sourceMap, err := d.createMap(source1)
	if err != nil {
		return err
	}
	logg.Log("sourceMap: %+v", sourceMap)

	// split the map into training and test
	trainMap, testMap, err := d.splitMap(sourceMap)
	if err != nil {
		return err
	}

	// iterate over the source
	logg.Log("iterate over source")
	for {
		hdr, err := source2.Next()

		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return err
		}

		// distribute to writers based on training and test maps
		var twToAdd *tar.Writer

		// if strings.HasPrefix(hdr.Name, "foo") {
		if trainMap.hasPath(hdr.Name) {
			// add to training tar writer
			twToAdd = train
		} else if testMap.hasPath(hdr.Name) {
			// add to testing tar writer
			twToAdd = test
		} else {
			logg.LogPanic("File not present in either test/train: %v", hdr.Name)
		}

		// TODO: is there a more efficient way to do this?
		bytes, err := ioutil.ReadAll(source2)
		if err != nil {
			return err
		}

		hdrToAdd := &tar.Header{
			Name: hdr.Name,
			Size: int64(len(bytes)),
		}
		if err := twToAdd.WriteHeader(hdrToAdd); err != nil {
			return err
		}
		if _, err := twToAdd.Write(bytes); err != nil {
			return err
		}

	}

	// close writers
	logg.LogTo("DATASET_SPLITTER", "Closing writers")
	if err := train.Close(); err != nil {
		return err
	}
	if err := test.Close(); err != nil {
		return err
	}
	logg.LogTo("DATASET_SPLITTER", "Closed writers")

	return nil
}

// Validate that the source tar stream conforms to expected specs
func (d DatasetSplitter) validate(source *tar.Reader) (bool, error) {

	// validation rules:
	// 1. has at least 2 files
	// 2. the depth of each file is 2 (folder/filename.xxx)

	numFiles := 0
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return false, err
		}
		numFiles += 1

		pathComponents := strings.Split(hdr.Name, "/")
		if len(pathComponents) != 2 {
			return false, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

	}

	if numFiles < 2 {
		return false, fmt.Errorf("Archive must contain at least 2 files")
	}

	return true, nil
}

// Create a map of folder -> []filename for all entries in the archive
func (d DatasetSplitter) createMap(source *tar.Reader) (filemap, error) {

	resultMap := filemap{}
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return nil, err
		}

		pathComponents := strings.Split(hdr.Name, "/")

		if len(pathComponents) != 2 {
			return nil, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

		directory := pathComponents[0]
		filename := pathComponents[1]

		resultMap.addFileToDirectory(directory, filename)

	}

	return resultMap, nil

}

// Split map into training and testing disjoint subsets based on values
// of DatasetSplitter's Dataset
func (d DatasetSplitter) splitMap(source filemap) (training filemap, testing filemap, err error) {

	training = filemap{}
	testing = filemap{}

	// iterate over source keys
	for directory, files := range source {

		numTraining := int(float64(len(files)) * d.Dataset.TrainingDataset.SplitPercentage)

		numTest := len(files) - int(numTraining)

		// split files into subsets based on ratios in dataset
		trainingFiles, testFiles, err := splitFiles(files, numTraining, numTest)

		if err != nil {
			return nil, nil, err
		}

		// add to respective maps
		training[directory] = trainingFiles
		testing[directory] = testFiles

	}

	return training, testing, nil

}

func splitFiles(files []string, numTraining, numTest int) (training []string, test []string, err error) {
	for i, file := range files {
		if i < numTraining {
			training = append(training, file)
		} else {
			test = append(test, file)
		}
	}
	return
}
<|endoftext|>