<commit_before>package google

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"

	"google.golang.org/api/storage/v1"
)

func resourceStorageBucketAcl() *schema.Resource {
	return &schema.Resource{
		Create:        resourceStorageBucketAclCreate,
		Read:          resourceStorageBucketAclRead,
		Update:        resourceStorageBucketAclUpdate,
		Delete:        resourceStorageBucketAclDelete,
		CustomizeDiff: resourceStorageRoleEntityCustomizeDiff,

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

			"default_acl": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"predefined_acl": {
				Type:          schema.TypeString,
				Optional:      true,
				ForceNew:      true,
				ConflictsWith: []string{"role_entity"},
			},

			"role_entity": {
				Type:          schema.TypeList,
				Optional:      true,
				Computed:      true,
				Elem:          &schema.Schema{Type: schema.TypeString},
				ConflictsWith: []string{"predefined_acl"},
			},
		},
	}
}

func resourceStorageRoleEntityCustomizeDiff(diff *schema.ResourceDiff, meta interface{}) error {
	keys := diff.GetChangedKeysPrefix("role_entity")
	if len(keys) < 1 {
		return nil
	}
	count := diff.Get("role_entity.#").(int)
	if count < 1 {
		return nil
	}
	state := map[string]struct{}{}
	conf := map[string]struct{}{}
	for i := 0; i < count; i++ {
		old, new := diff.GetChange(fmt.Sprintf("role_entity.%d", i))
		state[old.(string)] = struct{}{}
		conf[new.(string)] = struct{}{}
	}
	if len(state) != len(conf) {
		return nil
	}
	for k := range state {
		if _, ok := conf[k]; !ok {
			return nil
		}
	}
	return diff.Clear("role_entity")
}

type RoleEntity struct {
	Role   string
	Entity string
}

func getBucketAclId(bucket string) string {
	return bucket + "-acl"
}

func getRoleEntityPair(role_entity string) (*RoleEntity, error) {
	split := strings.Split(role_entity, ":")
	if len(split) != 2 {
		return nil, fmt.Errorf("Error, each role entity pair must be " +
			"formatted as ROLE:entity")
	}

	return &RoleEntity{Role: split[0], Entity: split[1]}, nil
}

func resourceStorageBucketAclCreate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	bucket := d.Get("bucket").(string)
	predefined_acl := ""
	default_acl := ""
	role_entity := make([]interface{}, 0)

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

	if v, ok := d.GetOk("role_entity"); ok {
		role_entity = v.([]interface{})
	}

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

	if len(predefined_acl) > 0 {
		res, err := config.clientStorage.Buckets.Get(bucket).Do()

		if err != nil {
			return fmt.Errorf("Error reading bucket %s: %v", bucket, err)
		}

		_, err = config.clientStorage.Buckets.Update(bucket,
			res).PredefinedAcl(predefined_acl).Do()

		if err != nil {
			return fmt.Errorf("Error updating bucket %s: %v", bucket, err)
		}

	}

	if len(role_entity) > 0 {
		current, err := config.clientStorage.BucketAccessControls.List(bucket).Do()
		if err != nil {
			return fmt.Errorf("Error retrieving current ACLs: %s", err)
		}
		for _, v := range role_entity {
			pair, err := getRoleEntityPair(v.(string))
			if err != nil {
				return err
			}
			var alreadyInserted bool
			for _, cur := range current.Items {
				if cur.Entity == pair.Entity && cur.Role == pair.Role {
					alreadyInserted = true
					break
				}
			}
			if alreadyInserted {
				log.Printf("[DEBUG]: pair %s-%s already exists, not trying to insert again\n", pair.Role, pair.Entity)
				continue
			}
			bucketAccessControl := &storage.BucketAccessControl{
				Role:   pair.Role,
				Entity: pair.Entity,
			}

			log.Printf("[DEBUG]: storing re %s-%s", pair.Role, pair.Entity)

			_, err = config.clientStorage.BucketAccessControls.Insert(bucket, bucketAccessControl).Do()

			if err != nil {
				return fmt.Errorf("Error updating ACL for bucket %s: %v", bucket, err)
			}
		}

	}

	if len(default_acl) > 0 {
		res, err := config.clientStorage.Buckets.Get(bucket).Do()

		if err != nil {
			return fmt.Errorf("Error reading bucket %s: %v", bucket, err)
		}

		_, err = config.clientStorage.Buckets.Update(bucket,
			res).PredefinedDefaultObjectAcl(default_acl).Do()

		if err != nil {
			return fmt.Errorf("Error updating bucket %s: %v", bucket, err)
		}

	}

	d.SetId(getBucketAclId(bucket))
	return resourceStorageBucketAclRead(d, meta)
}

func resourceStorageBucketAclRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	bucket := d.Get("bucket").(string)

	// The API offers no way to retrieve predefined ACLs,
	// and we can't tell which access controls were created
	// by the predefined roles, so...
	//
	// This is, needless to say, a bad state of affairs and
	// should be fixed.
	if _, ok := d.GetOk("role_entity"); ok {
		res, err := config.clientStorage.BucketAccessControls.List(bucket).Do()

		if err != nil {
			return handleNotFoundError(err, d, fmt.Sprintf("Storage Bucket ACL for bucket %q", d.Get("bucket").(string)))
		}
		entities := make([]string, 0, len(res.Items))
		for _, item := range res.Items {
			entities = append(entities, item.Role+":"+item.Entity)
		}

		d.Set("role_entity", entities)
	} else {
		// if we don't set `role_entity` to nil (effectively setting it
		// to empty in Terraform state), because it's computed now,
		// Terraform will think it's missing from state, is supposed
		// to be there, and throw up a diff for role_entity.#. So it
		// must always be set in state.
		d.Set("role_entity", nil)
	}

	return nil
}

func resourceStorageBucketAclUpdate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	bucket := d.Get("bucket").(string)

	if d.HasChange("role_entity") {
		bkt, err := config.clientStorage.Buckets.Get(bucket).Do()
		if err != nil {
			return fmt.Errorf("Error reading bucket %q: %v", bucket, err)
		}

		project := strconv.FormatUint(bkt.ProjectNumber, 10)
		o, n := d.GetChange("role_entity")
		old_re, new_re := o.([]interface{}), n.([]interface{})

		old_re_map := make(map[string]string)
		for _, v := range old_re {
			res, err := getRoleEntityPair(v.(string))

			if err != nil {
				return fmt.Errorf(
					"Old state has malformed Role/Entity pair: %v", err)
			}

			old_re_map[res.Entity] = res.Role
		}

		for _, v := range new_re {
			pair, err := getRoleEntityPair(v.(string))

			bucketAccessControl := &storage.BucketAccessControl{
				Role:   pair.Role,
				Entity: pair.Entity,
			}

			// If the old state is missing this entity, it needs to be inserted
			if _, ok := old_re_map[pair.Entity]; !ok {
				_, err = config.clientStorage.BucketAccessControls.Insert(
					bucket, bucketAccessControl).Do()
			}

			// Now we only store the keys that have to be removed
			delete(old_re_map, pair.Entity)

			if err != nil {
				return fmt.Errorf("Error updating ACL for bucket %s: %v", bucket, err)
			}
		}

		for entity, role := range old_re_map {
			if entity == fmt.Sprintf("project-owners-%s", project) && role == "OWNER" {
				log.Printf("[WARN]: Skipping %s-%s; not deleting owner ACL.", role, entity)
				continue
			}
			log.Printf("[DEBUG]: removing entity %s", entity)
			err := config.clientStorage.BucketAccessControls.Delete(bucket, entity).Do()

			if err != nil {
				return fmt.Errorf("Error updating ACL for bucket %s: %v", bucket, err)
			}
		}

		return resourceStorageBucketAclRead(d, meta)
	}

	if d.HasChange("default_acl") {
		default_acl := d.Get("default_acl").(string)

		res, err := config.clientStorage.Buckets.Get(bucket).Do()

		if err != nil {
			return fmt.Errorf("Error reading bucket %s: %v", bucket, err)
		}

		_, err = config.clientStorage.Buckets.Update(bucket,
			res).PredefinedDefaultObjectAcl(default_acl).Do()

		if err != nil {
			return fmt.Errorf("Error updating bucket %s: %v", bucket, err)
		}

		return resourceStorageBucketAclRead(d, meta)
	}

	return nil
}

func resourceStorageBucketAclDelete(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	bucket := d.Get("bucket").(string)

	bkt, err := config.clientStorage.Buckets.Get(bucket).Do()
	if err != nil {
		return fmt.Errorf("Error retrieving bucket %q: %v", bucket, err)
	}
	project := strconv.FormatUint(bkt.ProjectNumber, 10)

	re_local := d.Get("role_entity").([]interface{})
	for _, v := range re_local {
		res, err := getRoleEntityPair(v.(string))
		if err != nil {
			return err
		}

		if res.Entity == fmt.Sprintf("project-owners-%s", project) && res.Role == "OWNER" {
			log.Printf("[WARN]: Skipping %s-%s; not deleting owner ACL.", res.Role, res.Entity)
			continue
		}

		log.Printf("[DEBUG]: removing entity %s", res.Entity)

		err = config.clientStorage.BucketAccessControls.Delete(bucket, res.Entity).Do()

		if err != nil {
			return fmt.Errorf("Error deleting entity %s ACL: %s", res.Entity, err)
		}
	}

	return nil
}
<commit_msg>Avoid permadiff trying to remove project-owners (#3016)<commit_after>package google

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"

	"google.golang.org/api/storage/v1"
)

func resourceStorageBucketAcl() *schema.Resource {
	return &schema.Resource{
		Create:        resourceStorageBucketAclCreate,
		Read:          resourceStorageBucketAclRead,
		Update:        resourceStorageBucketAclUpdate,
		Delete:        resourceStorageBucketAclDelete,
		CustomizeDiff: resourceStorageRoleEntityCustomizeDiff,

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

			"default_acl": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"predefined_acl": {
				Type:          schema.TypeString,
				Optional:      true,
				ForceNew:      true,
				ConflictsWith: []string{"role_entity"},
			},

			"role_entity": {
				Type:          schema.TypeList,
				Optional:      true,
				Computed:      true,
				Elem:          &schema.Schema{Type: schema.TypeString},
				ConflictsWith: []string{"predefined_acl"},
			},
		},
	}
}

func resourceStorageRoleEntityCustomizeDiff(diff *schema.ResourceDiff, meta interface{}) error {
	keys := diff.GetChangedKeysPrefix("role_entity")
	if len(keys) < 1 {
		return nil
	}
	count := diff.Get("role_entity.#").(int)
	if count < 1 {
		return nil
	}
	state := map[string]struct{}{}
	conf := map[string]struct{}{}
	for i := 0; i < count; i++ {
		old, new := diff.GetChange(fmt.Sprintf("role_entity.%d", i))
		state[old.(string)] = struct{}{}
		conf[new.(string)] = struct{}{}
	}
	if len(state) != len(conf) {
		return nil
	}
	for k := range state {
		if _, ok := conf[k]; !ok {
			// project-owners- is explicitly stripped from the roles that this
			// resource will delete
			if strings.Contains(k, "OWNER:project-owners-") {
				continue
			}
			return nil
		}
	}
	return diff.Clear("role_entity")
}

type RoleEntity struct {
	Role   string
	Entity string
}

func getBucketAclId(bucket string) string {
	return bucket + "-acl"
}

func getRoleEntityPair(role_entity string) (*RoleEntity, error) {
	split := strings.Split(role_entity, ":")
	if len(split) != 2 {
		return nil, fmt.Errorf("Error, each role entity pair must be " +
			"formatted as ROLE:entity")
	}

	return &RoleEntity{Role: split[0], Entity: split[1]}, nil
}

func resourceStorageBucketAclCreate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	bucket := d.Get("bucket").(string)
	predefined_acl := ""
	default_acl := ""
	role_entity := make([]interface{}, 0)

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

	if v, ok := d.GetOk("role_entity"); ok {
		role_entity = v.([]interface{})
	}

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

	if len(predefined_acl) > 0 {
		res, err := config.clientStorage.Buckets.Get(bucket).Do()

		if err != nil {
			return fmt.Errorf("Error reading bucket %s: %v", bucket, err)
		}

		_, err = config.clientStorage.Buckets.Update(bucket,
			res).PredefinedAcl(predefined_acl).Do()

		if err != nil {
			return fmt.Errorf("Error updating bucket %s: %v", bucket, err)
		}

	}

	if len(role_entity) > 0 {
		current, err := config.clientStorage.BucketAccessControls.List(bucket).Do()
		if err != nil {
			return fmt.Errorf("Error retrieving current ACLs: %s", err)
		}
		for _, v := range role_entity {
			pair, err := getRoleEntityPair(v.(string))
			if err != nil {
				return err
			}
			var alreadyInserted bool
			for _, cur := range current.Items {
				if cur.Entity == pair.Entity && cur.Role == pair.Role {
					alreadyInserted = true
					break
				}
			}
			if alreadyInserted {
				log.Printf("[DEBUG]: pair %s-%s already exists, not trying to insert again\n", pair.Role, pair.Entity)
				continue
			}
			bucketAccessControl := &storage.BucketAccessControl{
				Role:   pair.Role,
				Entity: pair.Entity,
			}

			log.Printf("[DEBUG]: storing re %s-%s", pair.Role, pair.Entity)

			_, err = config.clientStorage.BucketAccessControls.Insert(bucket, bucketAccessControl).Do()

			if err != nil {
				return fmt.Errorf("Error updating ACL for bucket %s: %v", bucket, err)
			}
		}

	}

	if len(default_acl) > 0 {
		res, err := config.clientStorage.Buckets.Get(bucket).Do()

		if err != nil {
			return fmt.Errorf("Error reading bucket %s: %v", bucket, err)
		}

		_, err = config.clientStorage.Buckets.Update(bucket,
			res).PredefinedDefaultObjectAcl(default_acl).Do()

		if err != nil {
			return fmt.Errorf("Error updating bucket %s: %v", bucket, err)
		}

	}

	d.SetId(getBucketAclId(bucket))
	return resourceStorageBucketAclRead(d, meta)
}

func resourceStorageBucketAclRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	bucket := d.Get("bucket").(string)

	// The API offers no way to retrieve predefined ACLs,
	// and we can't tell which access controls were created
	// by the predefined roles, so...
	//
	// This is, needless to say, a bad state of affairs and
	// should be fixed.
	if _, ok := d.GetOk("role_entity"); ok {
		res, err := config.clientStorage.BucketAccessControls.List(bucket).Do()

		if err != nil {
			return handleNotFoundError(err, d, fmt.Sprintf("Storage Bucket ACL for bucket %q", d.Get("bucket").(string)))
		}
		entities := make([]string, 0, len(res.Items))
		for _, item := range res.Items {
			entities = append(entities, item.Role+":"+item.Entity)
		}

		d.Set("role_entity", entities)
	} else {
		// if we don't set `role_entity` to nil (effectively setting it
		// to empty in Terraform state), because it's computed now,
		// Terraform will think it's missing from state, is supposed
		// to be there, and throw up a diff for role_entity.#. So it
		// must always be set in state.
		d.Set("role_entity", nil)
	}

	return nil
}

func resourceStorageBucketAclUpdate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	bucket := d.Get("bucket").(string)

	if d.HasChange("role_entity") {
		bkt, err := config.clientStorage.Buckets.Get(bucket).Do()
		if err != nil {
			return fmt.Errorf("Error reading bucket %q: %v", bucket, err)
		}

		project := strconv.FormatUint(bkt.ProjectNumber, 10)
		o, n := d.GetChange("role_entity")
		old_re, new_re := o.([]interface{}), n.([]interface{})

		old_re_map := make(map[string]string)
		for _, v := range old_re {
			res, err := getRoleEntityPair(v.(string))

			if err != nil {
				return fmt.Errorf(
					"Old state has malformed Role/Entity pair: %v", err)
			}

			old_re_map[res.Entity] = res.Role
		}

		for _, v := range new_re {
			pair, err := getRoleEntityPair(v.(string))

			bucketAccessControl := &storage.BucketAccessControl{
				Role:   pair.Role,
				Entity: pair.Entity,
			}

			// If the old state is missing this entity, it needs to be inserted
			if _, ok := old_re_map[pair.Entity]; !ok {
				_, err = config.clientStorage.BucketAccessControls.Insert(
					bucket, bucketAccessControl).Do()
			}

			// Now we only store the keys that have to be removed
			delete(old_re_map, pair.Entity)

			if err != nil {
				return fmt.Errorf("Error updating ACL for bucket %s: %v", bucket, err)
			}
		}

		for entity, role := range old_re_map {
			if entity == fmt.Sprintf("project-owners-%s", project) && role == "OWNER" {
				log.Printf("[WARN]: Skipping %s-%s; not deleting owner ACL.", role, entity)
				continue
			}
			log.Printf("[DEBUG]: removing entity %s", entity)
			err := config.clientStorage.BucketAccessControls.Delete(bucket, entity).Do()

			if err != nil {
				return fmt.Errorf("Error updating ACL for bucket %s: %v", bucket, err)
			}
		}

		return resourceStorageBucketAclRead(d, meta)
	}

	if d.HasChange("default_acl") {
		default_acl := d.Get("default_acl").(string)

		res, err := config.clientStorage.Buckets.Get(bucket).Do()

		if err != nil {
			return fmt.Errorf("Error reading bucket %s: %v", bucket, err)
		}

		_, err = config.clientStorage.Buckets.Update(bucket,
			res).PredefinedDefaultObjectAcl(default_acl).Do()

		if err != nil {
			return fmt.Errorf("Error updating bucket %s: %v", bucket, err)
		}

		return resourceStorageBucketAclRead(d, meta)
	}

	return nil
}

func resourceStorageBucketAclDelete(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	bucket := d.Get("bucket").(string)

	bkt, err := config.clientStorage.Buckets.Get(bucket).Do()
	if err != nil {
		return fmt.Errorf("Error retrieving bucket %q: %v", bucket, err)
	}
	project := strconv.FormatUint(bkt.ProjectNumber, 10)

	re_local := d.Get("role_entity").([]interface{})
	for _, v := range re_local {
		res, err := getRoleEntityPair(v.(string))
		if err != nil {
			return err
		}

		if res.Entity == fmt.Sprintf("project-owners-%s", project) && res.Role == "OWNER" {
			log.Printf("[WARN]: Skipping %s-%s; not deleting owner ACL.", res.Role, res.Entity)
			continue
		}

		log.Printf("[DEBUG]: removing entity %s", res.Entity)

		err = config.clientStorage.BucketAccessControls.Delete(bucket, res.Entity).Do()

		if err != nil {
			return fmt.Errorf("Error deleting entity %s ACL: %s", res.Entity, err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package neoutils

import (
	"fmt"

	"github.com/jmcvetta/neoism"
)

type CypherRunner interface {
	CypherBatch(queries []*neoism.CypherQuery) error
}
<commit_msg>Added interfaces for common expected functionality<commit_after>package neoutils

import (
	"github.com/jmcvetta/neoism"
)

type CypherRunner interface {
	CypherBatch(queries []*neoism.CypherQuery) error
}

type IndexEnsurer interface {
	EnsureConstraints(indexes map[string]string) error
	EnsureIndexes(indexes map[string]string) error
}

type NeoConnection interface {
	CypherRunner
	IndexEnsurer
}
<|endoftext|>


<commit_before>package middlewares

import (
	"../config"
	"encoding/base64"
	"net/http"
	"strings"
)

func PlainText(h http.Handler) http.Handler {
	handler := func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(handler)
}

func Secure(h http.Handler) http.Handler {
	handler := func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization") // Basic dXJAKLRYHA
		if !strings.HasPrefix(auth, "Basic") || len(auth) == 0 {
			forceAuth(w)
			return
		}
		if len(auth) > 0 {
			pass, err := base64.StdEncoding.DecodeString(auth[6:])
			if err != nil || !strings.HasSuffix(string(pass), config.C.Pass) {
				forceAuth(w)
				return
			}
			h.ServeHTTP(w, r)
		}
	}
	return http.HandlerFunc(handler)
}

// HELPERS

func forceAuth(w http.ResponseWriter) {
	w.Header().Set("WWW-Authenticate", `Basic realm="TaskMgmt"`)
	w.WriteHeader(http.StatusUnauthorized)
	w.Write([]byte("Restricted!"))
}
<commit_msg>added json middleware<commit_after>package middlewares

import (
	"../config"
	"encoding/base64"
	"net/http"
	"strings"
)

func PlainText(h http.Handler) http.Handler {
	handler := func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(handler)
}

func JSON(h http.Handler) http.Handler {
	handler := func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(handler)
}

func Secure(h http.Handler) http.Handler {
	handler := func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization") // Basic dXJAKLRYHA
		if !strings.HasPrefix(auth, "Basic") || len(auth) == 0 {
			forceAuth(w)
			return
		}
		if len(auth) > 0 {
			pass, err := base64.StdEncoding.DecodeString(auth[6:])
			if err != nil || !strings.HasSuffix(string(pass), config.C.Pass) {
				forceAuth(w)
				return
			}
			h.ServeHTTP(w, r)
		}
	}
	return http.HandlerFunc(handler)
}

// HELPERS

func forceAuth(w http.ResponseWriter) {
	w.Header().Set("WWW-Authenticate", `Basic realm="TaskMgmt"`)
	w.WriteHeader(http.StatusUnauthorized)
	w.Write([]byte("Restricted!"))
}
<|endoftext|>


<commit_before>package aviator_test

import (
	"os"

	. "github.com/JulzDiverse/aviator/aviator"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Aviator", func() {

	var file string

	Context("Test basic function functionality", func() {

		BeforeEach(func() {
			file = `spruce:
  - base: base.yml
    merge:
    - with:
        files:
        - another.yml
    to: some/destination/file.yml

fly:
  config: pipeline.yml
  vars:
  - personal.yml`
		})

		Context("ReadYaml", func() {
			It("should return an non empty aviator struct", func() {
				Expect(ReadYaml([]byte(file))).ShouldNot(BeNil())
				Expect(len(ReadYaml([]byte(file)).Spruce)).To(Equal(1))
			})

			It("should contain a fly and a spruce object with params", func() {
				Expect(ReadYaml([]byte(file)).Fly.Config).To(Equal("pipeline.yml"))
				Expect(ReadYaml([]byte(file)).Spruce[0].Base).To(Equal("base.yml"))
			})
		})

		Context("ConcatFileName", func() {
			It("should concat a file name concatinaed by the parant folder name and the file name", func() {
				filePath := "path/to/some/file.yml"
				fn, _ := ConcatFileName(filePath)
				Expect(fn).To(Equal("some_file.yml"))
			})
		})

	})

	Context("Integration Tests: Spruce Specific Files", func() {

		BeforeEach(func() {
			file = `spruce:
- base: ../integration/yamls/base.yml
  merge:
  - with:
      files:
      - another.yml
      in_dir: ../integration/yamls/
  - with:
      files:
      - ../integration/yamls/addons/sub1/file1.yml
  to: ../integration/tmp/tmp.yml
- base: ../integration/tmp/tmp.yml
  merge:
  - with:
      files:
      - ../integration/yamls/yet-another.yml
  to: ../integration/tmp/result.yml`
		})

		Context("ProcessSpruceChain", func() {
			It("Should generate a result.yml file", func() {
				ProcessSprucePlan(ReadYaml([]byte(file)).Spruce)
				Expect("../integration/tmp/result.yml").To(BeAnExistingFile())
			})
		})

		Context("Cleanup", func() {
			It("Should delete all files in tmp", func() {
				Cleanup("../integration/tmp/")
				Expect("../integration/tmp/result.yml").ShouldNot(BeAnExistingFile())
			})
		})
	})

	Context("Resolve env vars", func() {
		var file, resultFile string
		BeforeEach(func() {
			file = `spruce:
  - base: $BASE
    merge:
    - with:
        files:
        - ${INPUT_FILE}
    to: $TMP_DIR/destination/$FILE

fly:
  config: pipeline.yml
  vars:
  - personal.yml`

			resultFile = `spruce:
  - base: base.yml
    merge:
    - with:
        files:
        - another.yml
    to: some/destination/file.yml

fly:
  config: pipeline.yml
  vars:
  - personal.yml`

		})

		It("should resolve env variables", func() {
			os.Setenv("BASE", "base.yml")
			os.Setenv("INPUT_FILE", "another.yml")
			os.Setenv("TMP_DIR", "some")
			os.Setenv("FILE", "file.yml")

			resolved := ResolveEnvVars([]byte(file))
			Expect(string(resolved)).To(Equal(resultFile))
		})

	})

})
<commit_msg>Fix broken tests<commit_after>package aviator_test

import (
	"os"

	. "github.com/JulzDiverse/aviator/aviator"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Aviator", func() {

	var file string

	Context("Test basic function functionality", func() {

		BeforeEach(func() {
			file = `spruce:
  - base: base.yml
    merge:
    - with:
        files:
        - another.yml
    to: some/destination/file.yml

fly:
  config: pipeline.yml
  vars:
  - personal.yml`
		})

		Context("ReadYaml", func() {
			It("should return an non empty aviator struct", func() {
				Expect(ReadYaml([]byte(file))).ShouldNot(BeNil())
				Expect(len(ReadYaml([]byte(file)).Spruce)).To(Equal(1))
			})

			It("should contain a fly and a spruce object with params", func() {
				Expect(ReadYaml([]byte(file)).Fly.Config).To(Equal("pipeline.yml"))
				Expect(ReadYaml([]byte(file)).Spruce[0].Base).To(Equal("base.yml"))
			})
		})

		Context("ConcatFileName", func() {
			It("should concat a file name concatinaed by the parant folder name and the file name", func() {
				filePath := "path/to/some/file.yml"
				fn, _ := ConcatFileName(filePath)
				Expect(fn).To(Equal("some_file.yml"))
			})
		})

	})

	Context("Integration Tests: Spruce Specific Files", func() {

		BeforeEach(func() {
			file = `spruce:
- base: ../integration/yamls/base.yml
  merge:
  - with:
      files:
      - another.yml
      in_dir: ../integration/yamls/
  - with:
      files:
      - ../integration/yamls/addons/sub1/file1.yml
  to: ../integration/tmp/tmp.yml
- base: ../integration/tmp/tmp.yml
  merge:
  - with:
      files:
      - ../integration/yamls/yet-another.yml
  to: ../integration/tmp/result.yml`
		})

		Context("ProcessSpruceChain", func() {
			It("Should generate a result.yml file", func() {
				ProcessSprucePlan(ReadYaml([]byte(file)).Spruce, false, false)
				Expect("../integration/tmp/result.yml").To(BeAnExistingFile())
			})
		})

		Context("Cleanup", func() {
			It("Should delete all files in tmp", func() {
				Cleanup("../integration/tmp/")
				Expect("../integration/tmp/result.yml").ShouldNot(BeAnExistingFile())
			})
		})
	})

	Context("Resolve env vars", func() {
		var file, resultFile string
		BeforeEach(func() {
			file = `spruce:
  - base: $BASE
    merge:
    - with:
        files:
        - ${INPUT_FILE}
    to: $TMP_DIR/destination/$FILE

fly:
  config: pipeline.yml
  vars:
  - personal.yml`

			resultFile = `spruce:
  - base: base.yml
    merge:
    - with:
        files:
        - another.yml
    to: some/destination/file.yml

fly:
  config: pipeline.yml
  vars:
  - personal.yml`

		})

		It("should resolve env variables", func() {
			os.Setenv("BASE", "base.yml")
			os.Setenv("INPUT_FILE", "another.yml")
			os.Setenv("TMP_DIR", "some")
			os.Setenv("FILE", "file.yml")

			resolved := ResolveEnvVars([]byte(file))
			Expect(string(resolved)).To(Equal(resultFile))
		})

	})

})
<|endoftext|>


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

package backend

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"testing"
	"time"
)

var (
	muTestBasePort sync.Mutex
	testBasePort   = 8080
)

func Test_StartServer(t *testing.T) {
	muTestBasePort.Lock()
	port := testBasePort
	testBasePort++
	muTestBasePort.Unlock()

	srv, err := StartServer(port)
	if err != nil {
		t.Fatal(err)
	}

	resp, err := http.Get(srv.addr + "/start")
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	fmt.Println(string(b))

	time.Sleep(time.Second)

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

package backend

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"testing"
	"time"
)

var (
	testMu       sync.Mutex
	testBasePort = 8080
)

func Test_StartServer(t *testing.T) {
	testMu.Lock()
	port := testBasePort
	testBasePort++
	testMu.Unlock()

	srv, err := StartServer(port)
	if err != nil {
		t.Fatal(err)
	}

	resp, err := http.Get(srv.addr + "/start")
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	fmt.Println(string(b))

	time.Sleep(time.Second)

	srv.Stop()
}
<|endoftext|>


<commit_before>package yaml

import (
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/fsnotify/fsnotify"
	"github.com/kelseyhightower/confd/log"
	"gopkg.in/yaml.v2"
)

var replacer = strings.NewReplacer("/", "_")

// Client provides a shell for the yaml client
type Client struct {
	filepath string
}

func NewYamlClient(filepath string) (*Client, error) {
	return &Client{filepath}, nil
}

func (c *Client) GetValues(keys []string) (map[string]string, error) {
	yamlMap := make(map[string]string)
	vars := make(map[string]string)

	data, err := ioutil.ReadFile(c.filepath)
	if err != nil {
		return vars, err
	}
	err = yaml.Unmarshal(data, &yamlMap)
	if err != nil {
		return vars, err
	}

	for _, key := range keys {
		k := transform(key)
		for yamlKey, yamlValue := range yamlMap {
			if strings.HasPrefix(yamlKey, k) {
				vars[clean(yamlKey)] = yamlValue
			}
		}
	}
	log.Debug(fmt.Sprintf("Key Map: %#v", vars))

	return vars, nil
}

func transform(key string) string {
	k := strings.TrimPrefix(key, "/")
	return replacer.Replace(k)
}

var cleanReplacer = strings.NewReplacer("_", "/")

func clean(key string) string {
	newKey := "/" + key
	return cleanReplacer.Replace(newKey)
}

func (c *Client) WatchPrefix(prefix string, waitIndex uint64, stopChan chan bool) (uint64, error) {
	if waitIndex == 0 {
		return 1, nil
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return 0, err
	}
	defer watcher.Close()

	err = watcher.Add(c.filepath)
	if err != nil {
		return 0, err
	}

	for {
		select {
		case event := <-watcher.Events:
			if event.Op&fsnotify.Write == fsnotify.Write {
				return 1, nil
			}
		case err := <-watcher.Errors:
			return 0, err
		case <-stopChan:
			return 0, nil
		}
	}
	return waitIndex, nil
}
<commit_msg>allow yaml backend to scope watches to specific keys<commit_after>package yaml

import (
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/fsnotify/fsnotify"
	"github.com/kelseyhightower/confd/log"
	"gopkg.in/yaml.v2"
)

var replacer = strings.NewReplacer("/", "_")

// Client provides a shell for the yaml client
type Client struct {
	filepath string
}

func NewYamlClient(filepath string) (*Client, error) {
	return &Client{filepath}, nil
}

func (c *Client) GetValues(keys []string) (map[string]string, error) {
	yamlMap := make(map[string]string)
	vars := make(map[string]string)

	data, err := ioutil.ReadFile(c.filepath)
	if err != nil {
		return vars, err
	}
	err = yaml.Unmarshal(data, &yamlMap)
	if err != nil {
		return vars, err
	}

	for _, key := range keys {
		k := transform(key)
		for yamlKey, yamlValue := range yamlMap {
			if strings.HasPrefix(yamlKey, k) {
				vars[clean(yamlKey)] = yamlValue
			}
		}
	}
	log.Debug(fmt.Sprintf("Key Map: %#v", vars))

	return vars, nil
}

func transform(key string) string {
	k := strings.TrimPrefix(key, "/")
	return replacer.Replace(k)
}

var cleanReplacer = strings.NewReplacer("_", "/")

func clean(key string) string {
	newKey := "/" + key
	return cleanReplacer.Replace(newKey)
}

func (c *Client) WatchPrefix(prefix string, keys []string, waitIndex uint64, stopChan chan bool) (uint64, error) {
	if waitIndex == 0 {
		return 1, nil
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return 0, err
	}
	defer watcher.Close()

	err = watcher.Add(c.filepath)
	if err != nil {
		return 0, err
	}

	for {
		select {
		case event := <-watcher.Events:
			if event.Op&fsnotify.Write == fsnotify.Write {
				return 1, nil
			}
		case err := <-watcher.Errors:
			return 0, err
		case <-stopChan:
			return 0, nil
		}
	}
	return waitIndex, nil
}
<|endoftext|>


<commit_before>package logging

import (
	"fmt"
	"log"
	"os"
	"path"
	"runtime"
	"strconv"
	"sync/atomic"
	"time"
)

const (
	logComMessage = iota
	logComClose
)

const (
	logInfoTypeError = iota
	logInfoTypeWarning
	logInfoTypeInfo
	logInfoTypeTrace
)

type logMessage struct {
	command int
	message string
}

//Logger is structure for internal use
type Logger struct {
	chanal      chan logMessage
	isOuted     chan struct{}
	file        *os.File
	canwrite    int32
	fileName    string
	maxFileSize int64
	level       byte
}

func addPrefix(infotype int, msg string, a ...interface{}) string {
	var m string
	switch infotype {
	case logInfoTypeError:
		m = " [E] " + msg
	case logInfoTypeWarning:
		m = " [W] " + msg
	case logInfoTypeTrace:
		m = " [T] " + msg
	default:
		m = " [I] " + msg
	}
	return fmt.Sprintf(m, a...)
}

func intToStrWithZero(Num, count int) string {
	const zeros string = "000000000000000000000000000000"
	str := strconv.Itoa(Num)
	return zeros[1:count-len(str)+1] + str
}

func (logger *Logger) startRotation() {
	for {
		//time.Sleep(100)
		data, ok := <-logger.chanal
		if !ok || data.command == logComClose {
			break
		}
		logger.saveMessage(data.message)
	}
	logger.isOuted <- struct{}{}
}

func (logger *Logger) saveHeader() {
	writeStrToFile(logger.file, "")
	writeStrToFile(logger.file, "----- Started -----")
	writeStrToFile(logger.file, fmt.Sprintf("NumCPU:[%d] OS:[%s] Arch:[%s]", runtime.NumCPU(), runtime.GOOS, runtime.GOARCH))
	writeStrToFile(logger.file, "")
	writeStrToFile(logger.file, "-------------------")
}

func writeStrToFile(file *os.File, message string) error {
	str := time.Now().Format("2006-01-02 15:04:05.000") + message + "\n"
	buffer := []byte(str)
	if _, err := file.Write(buffer); err != nil {
		return err
	}
	return nil
}

func (logger *Logger) saveMessage(msg string) error {
	pos, err := logger.file.Seek(0, 2)
	if err != nil {
		return err
	}
	if pos > logger.maxFileSize {
		logger.file.Close()
		ext := path.Ext(logger.fileName)
		bakfile := logger.fileName[0:len(logger.fileName)-len(ext)] + ".bak"
		if err = os.Remove(bakfile); err == nil || os.IsNotExist(err) {
			os.Rename(logger.fileName, bakfile)
		}
		logger.file, err = os.OpenFile(logger.fileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		if err != nil {
			atomic.StoreInt32(&logger.canwrite, 0)
			return err
		}
		logger.saveHeader()
	}

	if err = writeStrToFile(logger.file, msg); err != nil {
		atomic.StoreInt32(&logger.canwrite, 0)
		return err
	}
	return nil
}

//CreateLog is created new logging system
func CreateLog(fileName string, maxFileSize int64, level byte) (logger *Logger, err error) {
	logger = new(Logger)
	logger.fileName = fileName
	logger.maxFileSize = maxFileSize
	logger.level = level
	logger.canwrite = 0
	logger.chanal = make(chan logMessage, 1000)
	logger.isOuted = make(chan struct{})
	err = nil
	logger.file, err = os.OpenFile(logger.fileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		return nil, err
	}
	atomic.StoreInt32(&logger.canwrite, 1)
	logger.saveHeader()
	go logger.startRotation()
	return
}

func (logger *Logger) infoout(infotype, id int, msg string, a ...interface{}) {
	m := addPrefix(infotype, msg, a...)
	if logger == nil {
		log.Println(m)
		return
	}
	if atomic.AddInt32(&logger.canwrite, 0) == 0 {
		return
	}
	logger.chanal <- logMessage{
		command: logComMessage,
		message: m,
	}
}

// Trace outputs the message into log with Trace level
func (logger *Logger) Trace(msg string, a ...interface{}) {
	logger.infoout(logInfoTypeTrace, 0, msg, a...)
}

// Warning outputs the message into log with Warning level
func (logger *Logger) Warning(msg string, a ...interface{}) {
	logger.infoout(logInfoTypeWarning, 0, msg, a...)
}

// Error outputs the message into log with Warning level
func (logger *Logger) Error(msg string, a ...interface{}) {
	logger.infoout(logInfoTypeError, 0, msg, a...)
}

// Info outputs the message into log with Warning level
func (logger *Logger) Info(msg string, a ...interface{}) {
	logger.infoout(logInfoTypeInfo, 0, msg, a...)
}

// Close closes all opened handles and stop logging
func (logger *Logger) Close() {

	logger.chanal <- logMessage{
		command: logComClose,
		message: "",
	}
	<-logger.isOuted
	logger.file.Close()

}
<commit_msg>Fix some warnings from gometalinter list<commit_after>package logging

import (
	"fmt"
	"log"
	"os"
	"path"
	"runtime"
	"strconv"
	"sync/atomic"
	"time"
)

const (
	logComMessage = iota
	logComClose
)

const (
	logInfoTypeError = iota
	logInfoTypeWarning
	logInfoTypeInfo
	logInfoTypeTrace
)

type logMessage struct {
	command int
	message string
}

//Logger is structure for internal use
type Logger struct {
	chanal      chan logMessage
	isOuted     chan struct{}
	file        *os.File
	canwrite    int32
	fileName    string
	maxFileSize int64
	level       byte
}

func addPrefix(infotype int, msg string, a ...interface{}) string {
	var m string
	switch infotype {
	case logInfoTypeError:
		m = " [E] " + msg
	case logInfoTypeWarning:
		m = " [W] " + msg
	case logInfoTypeTrace:
		m = " [T] " + msg
	default:
		m = " [I] " + msg
	}
	return fmt.Sprintf(m, a...)
}

func intToStrWithZero(Num, count int) string {
	const zeros string = "000000000000000000000000000000"
	str := strconv.Itoa(Num)
	return zeros[1:count-len(str)+1] + str
}

func (logger *Logger) startRotation() {
	for {
		//time.Sleep(100)
		data, ok := <-logger.chanal
		if !ok || data.command == logComClose {
			break
		}
		logger.saveMessage(data.message)
	}
	logger.isOuted <- struct{}{}
}

func (logger *Logger) saveHeader() {
	writeStrToFile(logger.file, "")
	writeStrToFile(logger.file, "----- Started -----")
	writeStrToFile(logger.file, fmt.Sprintf("NumCPU:[%d] OS:[%s] Arch:[%s]", runtime.NumCPU(), runtime.GOOS, runtime.GOARCH))
	writeStrToFile(logger.file, "")
	writeStrToFile(logger.file, "-------------------")
}

func writeStrToFile(file *os.File, message string) error {
	str := time.Now().Format("2006-01-02 15:04:05.000") + message + "\n"
	buffer := []byte(str)
	if _, err := file.Write(buffer); err != nil {
		return err
	}
	return nil
}

func (logger *Logger) saveMessage(msg string) error {
	pos, err := logger.file.Seek(0, 2)
	if err != nil {
		return err
	}
	if pos > logger.maxFileSize {
		logger.file.Close()
		ext := path.Ext(logger.fileName)
		bakfile := logger.fileName[0:len(logger.fileName)-len(ext)] + ".bak"
		if err = os.Remove(bakfile); err == nil || os.IsNotExist(err) {
			os.Rename(logger.fileName, bakfile)
		}
		logger.file, err = os.OpenFile(logger.fileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		if err != nil {
			atomic.StoreInt32(&logger.canwrite, 0)
			return err
		}
		logger.saveHeader()
	}

	if err = writeStrToFile(logger.file, msg); err != nil {
		atomic.StoreInt32(&logger.canwrite, 0)
		return err
	}
	return nil
}

//CreateLog is created new logging system
func CreateLog(fileName string, maxFileSize int64, level byte) (logger *Logger, err error) {
	logger = new(Logger)
	logger.fileName = fileName
	logger.maxFileSize = maxFileSize
	logger.level = level
	logger.canwrite = 0
	logger.chanal = make(chan logMessage, 1000)
	logger.isOuted = make(chan struct{})
	err = nil
	logger.file, err = os.OpenFile(logger.fileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		return nil, err
	}
	atomic.StoreInt32(&logger.canwrite, 1)
	logger.saveHeader()
	go logger.startRotation()
	return
}

func (logger *Logger) infoout(infotype, id int, msg string, a ...interface{}) {
	m := addPrefix(infotype, msg, a...)
	if logger == nil {
		log.Println(m)
		return
	}
	if atomic.AddInt32(&logger.canwrite, 0) == 0 {
		return
	}
	logger.chanal <- logMessage{
		command: logComMessage,
		message: m,
	}
}

// Trace outputs the message into log with Trace level
func (logger *Logger) Trace(msg string, a ...interface{}) {
	logger.infoout(logInfoTypeTrace, 0, msg, a...)
}

// Warning outputs the message into log with Warning level
func (logger *Logger) Warning(msg string, a ...interface{}) {
	logger.infoout(logInfoTypeWarning, 0, msg, a...)
}

// Error outputs the message into log with Warning level
func (logger *Logger) Error(msg string, a ...interface{}) {
	logger.infoout(logInfoTypeError, 0, msg, a...)
}

// Info outputs the message into log with Warning level
func (logger *Logger) Info(msg string, a ...interface{}) {
	logger.infoout(logInfoTypeInfo, 0, msg, a...)
}

// Close closes all opened handles and stop logging
func (logger *Logger) Close() {

	logger.chanal <- logMessage{
		command: logComClose,
		message: "",
	}
	<-logger.isOuted
	logger.file.Close()

}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"

	"github.com/weaveworks/flux"
)

type setConfigOpts struct {
	*rootOpts
	file string
}

func newSetConfig(parent *rootOpts) *setConfigOpts {
	return &setConfigOpts{rootOpts: parent}
}

func (opts *setConfigOpts) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "set-config",
		Short: "set configuration values for an instance",
		Example: makeExample(
			"fluxctl config --file=./dev/flux-conf.yaml",
		),
		RunE: opts.RunE,
	}
	cmd.Flags().StringVarP(&opts.file, "file", "f", "", "A file to upload as configuration; this will overwrite all values.")
	return cmd
}

func (opts *setConfigOpts) RunE(_ *cobra.Command, args []string) error {
	if len(args) > 0 {
		return errorWantedNoArgs
	}

	if opts.file == "" {
		return newUsageError("-f, --file is required")
	}

	var config flux.InstanceConfig

	bytes, err := ioutil.ReadFile(opts.file)
	if err == nil {
		err = yaml.Unmarshal(bytes, &config)
	}
	if err != nil {
		return errors.Wrapf(err, "reading config from file")
	}

	return opts.API.SetConfig(noInstanceID, config)
}
<commit_msg>Correct name of subcommand in set-config example<commit_after>package main

import (
	"io/ioutil"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"

	"github.com/weaveworks/flux"
)

type setConfigOpts struct {
	*rootOpts
	file string
}

func newSetConfig(parent *rootOpts) *setConfigOpts {
	return &setConfigOpts{rootOpts: parent}
}

func (opts *setConfigOpts) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "set-config",
		Short: "set configuration values for an instance",
		Example: makeExample(
			"fluxctl set-config --file=./dev/flux-conf.yaml",
		),
		RunE: opts.RunE,
	}
	cmd.Flags().StringVarP(&opts.file, "file", "f", "", "A file to upload as configuration; this will overwrite all values.")
	return cmd
}

func (opts *setConfigOpts) RunE(_ *cobra.Command, args []string) error {
	if len(args) > 0 {
		return errorWantedNoArgs
	}

	if opts.file == "" {
		return newUsageError("-f, --file is required")
	}

	var config flux.InstanceConfig

	bytes, err := ioutil.ReadFile(opts.file)
	if err == nil {
		err = yaml.Unmarshal(bytes, &config)
	}
	if err != nil {
		return errors.Wrapf(err, "reading config from file")
	}

	return opts.API.SetConfig(noInstanceID, config)
}
<|endoftext|>


<commit_before>// Build sets scollector version information and should be run from the
// scollector directory.
package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

func main() {
	const path = "main.go"
	var hash, id string
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}
	for _, d := range f.Decls {
		switch d := d.(type) {
		case *ast.GenDecl:
			if d.Tok != token.CONST {
				continue
			}
			for _, spec := range d.Specs {
				switch spec := spec.(type) {
				case *ast.ValueSpec:
					if len(spec.Names) != 1 || len(spec.Values) != 1 {
						continue
					}
					switch spec.Names[0].Name {
					case "VersionDate":
						switch value := spec.Values[0].(type) {
						case *ast.BasicLit:
							id = time.Now().UTC().Format("20060102150405")
							value.Value = id
						}
					case "VersionID":
						switch value := spec.Values[0].(type) {
						case *ast.BasicLit:
							rev, err := exec.Command("git", "rev-parse", "HEAD").Output()
							if err != nil {
								log.Fatal(err)
							}
							hash = fmt.Sprintf(`"%s"`, strings.TrimSpace(string(rev)))
							value.Value = hash
						}
					}
				}
			}
		}
	}
	file, err := os.Open(path)
	if err != nil {
		log.Fatal(err)
	}
	info, err := file.Stat()
	if err != nil {
		log.Fatal(err)
	}
	var buf bytes.Buffer
	if err := printer.Fprint(&buf, fset, f); err != nil {
		log.Fatal(err)
	}
	fb, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(err)
	}
	if err := ioutil.WriteFile(path, fb, info.Mode()); err != nil {
		log.Fatal(err)
	}
	fmt.Printf("version:\n  hash: %s\n  id: %s\n", hash, id)
}
<commit_msg>cmd/scollector: Merge pull request #134 from Wessie/build-rewrite<commit_after>// Build sets scollector version information and should be run from the
// scollector directory.
package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

func main() {
	const path = "main.go"
	mainfile, err := os.OpenFile(path, os.O_RDWR, 0660)
	if err != nil {
		log.Fatal(err)
	}
	defer mainfile.Close()

	var hash, id string
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, mainfile, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}

	for _, d := range f.Decls {
		d, ok := d.(*ast.GenDecl)
		if !ok {
			continue
		}

		if d.Tok != token.CONST {
			continue
		}
		for _, spec := range d.Specs {
			spec, ok := spec.(*ast.ValueSpec)
			if !ok {
				continue
			}

			if len(spec.Names) != 1 || len(spec.Values) != 1 {
				continue
			}

			value, ok := spec.Values[0].(*ast.BasicLit)
			if !ok {
				continue
			}

			switch spec.Names[0].Name {
			case "VersionDate":
				id = time.Now().UTC().Format("20060102150405")
				value.Value = id
			case "VersionID":
				rev, err := exec.Command("git", "rev-parse", "HEAD").Output()
				if err != nil {
					log.Fatal(err)
				}
				hash = fmt.Sprintf(`"%s"`, strings.TrimSpace(string(rev)))
				value.Value = hash
			}
		}
	}

	var config = printer.Config{Mode: printer.UseSpaces | printer.TabIndent, Tabwidth: 8}
	var buf bytes.Buffer
	if err := config.Fprint(&buf, fset, f); err != nil {
		log.Fatal(err)
	}

	if _, err := mainfile.Seek(0, os.SEEK_SET); err != nil {
		log.Fatal(err)
	}

	if _, err := io.Copy(mainfile, &buf); err != nil {
		log.Fatal(err)
	}

	fmt.Printf("version:\n  hash: %s\n  id: %s\n", hash, id)
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"sync"

	"github.com/eahydra/socks"
)

type UpstreamDialer struct {
	forwardDialers []socks.Dialer
	nextRouter     int
	lock           sync.Mutex
}

func NewUpstreamDialer(forwardDialers []socks.Dialer) *UpstreamDialer {
	return &UpstreamDialer{
		forwardDialers: forwardDialers,
	}
}

func (u *UpstreamDialer) getNextDialer() socks.Dialer {
	u.lock.Lock()
	defer u.lock.Unlock()
	index := u.nextRouter
	u.nextRouter++
	if u.nextRouter >= len(u.forwardDialers) {
		u.nextRouter = 0
	}
	if index < len(u.forwardDialers) {
		return u.forwardDialers[index]
	}
	panic("unreached")
}

func (u *UpstreamDialer) Dial(network, address string) (net.Conn, error) {
	router := u.getNextDialer()
	conn, err := router.Dial(network, address)
	if err != nil {
		ErrLog.Println("UpstreamDialer router.Dial failed, err:", err, network, address)
		return nil, err
	}
	return conn, nil
}
<commit_msg>modify roundrobin implemention<commit_after>package main

import (
	"net"
	"sync/atomic"

	"github.com/eahydra/socks"
)

type UpstreamDialer struct {
	nextRouter     uint64
	forwardDialers []socks.Dialer
}

func NewUpstreamDialer(forwardDialers []socks.Dialer) *UpstreamDialer {
	return &UpstreamDialer{
		forwardDialers: forwardDialers,
	}
}

func (u *UpstreamDialer) getNextDialer() socks.Dialer {
	old := atomic.LoadUint64(&u.nextRouter)
	return u.forwardDialers[old%uint64(len(u.forwardDialers))]
}

func (u *UpstreamDialer) Dial(network, address string) (net.Conn, error) {
	router := u.getNextDialer()
	conn, err := router.Dial(network, address)
	if err != nil {
		ErrLog.Println("UpstreamDialer router.Dial failed, err:", err, network, address)
		return nil, err
	}
	return conn, nil
}
<|endoftext|>


<commit_before>package service

import (
	"errors"
	"sync"

	msg "github.com/jbenet/go-ipfs/net/message"
	u "github.com/jbenet/go-ipfs/util"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
)

// Handler is an interface that objects must implement in order to handle
// a service's requests.
type Handler interface {

	// HandleMessage receives an incoming message, and potentially returns
	// a response message to send back.
	HandleMessage(context.Context, msg.NetMessage) (msg.NetMessage, error)
}

// Service is a networking component that protocols can use to multiplex
// messages over the same channel, and to issue + handle requests.
type Service struct {
	// Handler is the object registered to handle incoming requests.
	Handler Handler

	// Requests are all the pending requests on this service.
	Requests     RequestMap
	RequestsLock sync.RWMutex

	// cancel is the function to stop the Service
	cancel context.CancelFunc

	// Message Pipe (connected to the outside world)
	*msg.Pipe
}

// NewService creates a service object with given type ID and Handler
func NewService(h Handler) *Service {
	return &Service{
		Handler:  h,
		Requests: RequestMap{},
		Pipe:     msg.NewPipe(10),
	}
}

// Start kicks off the Service goroutines.
func (s *Service) Start(ctx context.Context) error {
	if s.cancel != nil {
		return errors.New("Service already started.")
	}

	// make a cancellable context.
	ctx, s.cancel = context.WithCancel(ctx)

	go s.handleIncomingMessages(ctx)
	return nil
}

// Stop stops Service activity.
func (s *Service) Stop() {
	s.cancel()
	s.cancel = context.CancelFunc(nil)
}

// GetPipe implements the mux.Protocol interface
func (s *Service) GetPipe() *msg.Pipe {
	return s.Pipe
}

// sendMessage sends a message out (actual leg work. SendMessage is to export w/o rid)
func (s *Service) sendMessage(ctx context.Context, m msg.NetMessage, rid RequestID) error {

	// serialize ServiceMessage wrapper
	data, err := wrapData(m.Data(), rid)
	if err != nil {
		return err
	}

	// send message
	m2 := msg.New(m.Peer(), data)
	select {
	case s.Outgoing <- m2:
	case <-ctx.Done():
		return ctx.Err()
	}

	return nil
}

// SendMessage sends a message out
func (s *Service) SendMessage(ctx context.Context, m msg.NetMessage) error {
	return s.sendMessage(ctx, m, nil)
}

// SendRequest sends a request message out and awaits a response.
func (s *Service) SendRequest(ctx context.Context, m msg.NetMessage) (msg.NetMessage, error) {

	// create a request
	r, err := NewRequest(m.Peer().ID)
	if err != nil {
		return nil, err
	}

	// register Request
	s.RequestsLock.Lock()
	s.Requests[r.Key()] = r
	s.RequestsLock.Unlock()

	// defer deleting this request
	defer func() {
		s.RequestsLock.Lock()
		delete(s.Requests, r.Key())
		s.RequestsLock.Unlock()
	}()

	// check if we should bail after waiting for mutex
	select {
	default:
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// Send message
	s.sendMessage(ctx, m, r.ID)

	// wait for response
	m = nil
	err = nil
	select {
	case m = <-r.Response:
	case <-ctx.Done():
		err = ctx.Err()
	}

	return m, err
}

// handleIncoming consumes the messages on the s.Incoming channel and
// routes them appropriately (to requests, or handler).
func (s *Service) handleIncomingMessages(ctx context.Context) {
	for {
		select {
		case m := <-s.Incoming:
			go s.handleIncomingMessage(ctx, m)

		case <-ctx.Done():
			return
		}
	}
}

func (s *Service) handleIncomingMessage(ctx context.Context, m msg.NetMessage) {

	// unwrap the incoming message
	data, rid, err := unwrapData(m.Data())
	if err != nil {
		u.PErr("de-serializing error: %v\n", err)
	}
	m2 := msg.New(m.Peer(), data)

	// if it's a request (or has no RequestID), handle it
	if rid == nil || rid.IsRequest() {
		if s.Handler == nil {
			u.PErr("service dropped msg: %v\n", m)
			return // no handler, drop it.
		}

		// should this be "go HandleMessage ... ?"
		r1, err := s.Handler.HandleMessage(ctx, m2)
		if err != nil {
			u.PErr("handled message yielded error %v\n", err)
			return
		}

		// if handler gave us a response, send it back out!
		if r1 != nil {
			err := s.sendMessage(ctx, r1, rid.Response())
			if err != nil {
				u.PErr("error sending response message: %v\n", err)
			}
		}
		return
	}

	// Otherwise, it is a response. handle it.
	if !rid.IsResponse() {
		u.PErr("RequestID should identify a response here.\n")
	}

	key := RequestKey(m.Peer().ID, RequestID(rid))
	s.RequestsLock.RLock()
	r, found := s.Requests[key]
	s.RequestsLock.RUnlock()

	if !found {
		u.PErr("no request key %v (timeout?)\n", []byte(key))
		return
	}

	select {
	case r.Response <- m2:
	case <-ctx.Done():
	}
}

func (s *Service) SetHandler(h Handler) {
	s.Handler = h
}
<commit_msg>NoResponse service<commit_after>package service

import (
	"errors"
	"sync"

	msg "github.com/jbenet/go-ipfs/net/message"
	u "github.com/jbenet/go-ipfs/util"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
)

// ErrNoResponse is returned by Service when a Request did not get a response,
// and no other error happened
var ErrNoResponse = errors.New("no response to request")

// Handler is an interface that objects must implement in order to handle
// a service's requests.
type Handler interface {

	// HandleMessage receives an incoming message, and potentially returns
	// a response message to send back.
	HandleMessage(context.Context, msg.NetMessage) (msg.NetMessage, error)
}

// Service is a networking component that protocols can use to multiplex
// messages over the same channel, and to issue + handle requests.
type Service struct {
	// Handler is the object registered to handle incoming requests.
	Handler Handler

	// Requests are all the pending requests on this service.
	Requests     RequestMap
	RequestsLock sync.RWMutex

	// cancel is the function to stop the Service
	cancel context.CancelFunc

	// Message Pipe (connected to the outside world)
	*msg.Pipe
}

// NewService creates a service object with given type ID and Handler
func NewService(h Handler) *Service {
	return &Service{
		Handler:  h,
		Requests: RequestMap{},
		Pipe:     msg.NewPipe(10),
	}
}

// Start kicks off the Service goroutines.
func (s *Service) Start(ctx context.Context) error {
	if s.cancel != nil {
		return errors.New("Service already started.")
	}

	// make a cancellable context.
	ctx, s.cancel = context.WithCancel(ctx)

	go s.handleIncomingMessages(ctx)
	return nil
}

// Stop stops Service activity.
func (s *Service) Stop() {
	s.cancel()
	s.cancel = context.CancelFunc(nil)
}

// GetPipe implements the mux.Protocol interface
func (s *Service) GetPipe() *msg.Pipe {
	return s.Pipe
}

// sendMessage sends a message out (actual leg work. SendMessage is to export w/o rid)
func (s *Service) sendMessage(ctx context.Context, m msg.NetMessage, rid RequestID) error {

	// serialize ServiceMessage wrapper
	data, err := wrapData(m.Data(), rid)
	if err != nil {
		return err
	}

	// send message
	m2 := msg.New(m.Peer(), data)
	select {
	case s.Outgoing <- m2:
	case <-ctx.Done():
		return ctx.Err()
	}

	return nil
}

// SendMessage sends a message out
func (s *Service) SendMessage(ctx context.Context, m msg.NetMessage) error {
	return s.sendMessage(ctx, m, nil)
}

// SendRequest sends a request message out and awaits a response.
func (s *Service) SendRequest(ctx context.Context, m msg.NetMessage) (msg.NetMessage, error) {

	// create a request
	r, err := NewRequest(m.Peer().ID)
	if err != nil {
		return nil, err
	}

	// register Request
	s.RequestsLock.Lock()
	s.Requests[r.Key()] = r
	s.RequestsLock.Unlock()

	// defer deleting this request
	defer func() {
		s.RequestsLock.Lock()
		delete(s.Requests, r.Key())
		s.RequestsLock.Unlock()
	}()

	// check if we should bail after waiting for mutex
	select {
	default:
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// Send message
	s.sendMessage(ctx, m, r.ID)

	// wait for response
	m = nil
	err = nil
	select {
	case m = <-r.Response:
	case <-ctx.Done():
		err = ctx.Err()
	}

	if m == nil {
		return nil, ErrNoResponse
	}

	return m, err
}

// handleIncoming consumes the messages on the s.Incoming channel and
// routes them appropriately (to requests, or handler).
func (s *Service) handleIncomingMessages(ctx context.Context) {
	for {
		select {
		case m := <-s.Incoming:
			go s.handleIncomingMessage(ctx, m)

		case <-ctx.Done():
			return
		}
	}
}

func (s *Service) handleIncomingMessage(ctx context.Context, m msg.NetMessage) {

	// unwrap the incoming message
	data, rid, err := unwrapData(m.Data())
	if err != nil {
		u.PErr("de-serializing error: %v\n", err)
	}
	m2 := msg.New(m.Peer(), data)

	// if it's a request (or has no RequestID), handle it
	if rid == nil || rid.IsRequest() {
		if s.Handler == nil {
			u.PErr("service dropped msg: %v\n", m)
			return // no handler, drop it.
		}

		// should this be "go HandleMessage ... ?"
		r1, err := s.Handler.HandleMessage(ctx, m2)
		if err != nil {
			u.PErr("handled message yielded error %v\n", err)
			return
		}

		// if handler gave us a response, send it back out!
		if r1 != nil {
			err := s.sendMessage(ctx, r1, rid.Response())
			if err != nil {
				u.PErr("error sending response message: %v\n", err)
			}
		}
		return
	}

	// Otherwise, it is a response. handle it.
	if !rid.IsResponse() {
		u.PErr("RequestID should identify a response here.\n")
	}

	key := RequestKey(m.Peer().ID, RequestID(rid))
	s.RequestsLock.RLock()
	r, found := s.Requests[key]
	s.RequestsLock.RUnlock()

	if !found {
		u.PErr("no request key %v (timeout?)\n", []byte(key))
		return
	}

	select {
	case r.Response <- m2:
	case <-ctx.Done():
	}
}

// SetHandler assigns the request Handler for this service.
func (s *Service) SetHandler(h Handler) {
	s.Handler = h
}
<|endoftext|>


<commit_before>package urlutil

import (
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strings"
)

// ToSlug creates a slug byte array from an input byte array.
// Slugs have words separated by a hyphen with no punctuation
// or spaces.
func ToSlug(slug []byte) []byte {
	// Convert punctuation and spaces to hyphens: string([]byte{45}) = "-"
	slug = regexp.MustCompile(`[\*\s]+`).ReplaceAll(slug, []byte{45})
	return regexp.MustCompile(`(^-+|-+$)`).ReplaceAll(slug, []byte{})
}

// ToSlugLowerString creates a lower-cased slug string
func ToSlugLowerString(s string) string {
	return string(ToSlug([]byte(strings.ToLower(s))))
}

// BuildURLFromMap returns a URL as a string from a base URL and a
// set of query parameters as a map[string]string{}
func BuildURLFromMap(baseUrl string, queryParams map[string]string) string {
	if len(queryParams) < 1 {
		return baseUrl
	}
	queryValues := url.Values{}
	for key, val := range queryParams {
		queryValues.Set(key, val)
	}
	return BuildURL(baseUrl, queryValues)
}

// BuildURL returns a URL string from a base URL and url.Values.
func BuildURL(baseUrl string, queryValues url.Values) string {
	qryString := queryValues.Encode()
	if len(qryString) > 0 {
		return baseUrl + "?" + qryString
	}
	return baseUrl
}

// GetURLBody returns an HTTP response byte array body from
// a URL.
func GetURLBody(absoluteUrl string) ([]byte, error) {
	req, err := http.NewRequest("GET", absoluteUrl, nil)
	cli := &http.Client{}
	res, err := cli.Do(req)
	if err != nil {
		return []byte{}, err
	}
	defer res.Body.Close()
	return ioutil.ReadAll(res.Body)
}

// GetURLPostBody returns a HTTP post body as a byte array from a
// URL, body type and an io.Reader.
func GetURLPostBody(absoluteUrl string, bodyType string, reqBody io.Reader) ([]byte, error) {
	client := &http.Client{}
	res, err := client.Post(absoluteUrl, bodyType, reqBody)
	if err != nil {
		return []byte{}, err
	}
	defer res.Body.Close()
	return ioutil.ReadAll(res.Body)
}

// JoinAbsolute performs a path.Join() while preserving two slashes after the scheme.
func JoinAbsolute(elem ...string) string {
	return regexp.MustCompile(`^([A-Za-z]+:/)`).ReplaceAllString(path.Join(elem...), "${1}/")
}
<commit_msg>update ToSlug format<commit_after>package urlutil

import (
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strings"
)

// ToSlug creates a slug byte array from an input byte array.
// Slugs have words separated by a hyphen with no punctuation
// or spaces.
func ToSlug(slug []byte) []byte {
	// Convert punctuation and spaces to hyphens: string([]byte{45}) = "-"
	slug = regexp.MustCompile(`[\*\s]+`).ReplaceAll(slug, []byte{45})
	slug = regexp.MustCompile(`["']+`).ReplaceAll(slug, []byte{})
	return regexp.MustCompile(`(^-+|-+$)`).ReplaceAll(slug, []byte{})
}

// ToSlugLowerString creates a lower-cased slug string
func ToSlugLowerString(s string) string {
	return string(ToSlug([]byte(strings.ToLower(s))))
}

// BuildURLFromMap returns a URL as a string from a base URL and a
// set of query parameters as a map[string]string{}
func BuildURLFromMap(baseUrl string, queryParams map[string]string) string {
	if len(queryParams) < 1 {
		return baseUrl
	}
	queryValues := url.Values{}
	for key, val := range queryParams {
		queryValues.Set(key, val)
	}
	return BuildURL(baseUrl, queryValues)
}

// BuildURL returns a URL string from a base URL and url.Values.
func BuildURL(baseUrl string, queryValues url.Values) string {
	qryString := queryValues.Encode()
	if len(qryString) > 0 {
		return baseUrl + "?" + qryString
	}
	return baseUrl
}

// GetURLBody returns an HTTP response byte array body from
// a URL.
func GetURLBody(absoluteUrl string) ([]byte, error) {
	req, err := http.NewRequest("GET", absoluteUrl, nil)
	cli := &http.Client{}
	res, err := cli.Do(req)
	if err != nil {
		return []byte{}, err
	}
	defer res.Body.Close()
	return ioutil.ReadAll(res.Body)
}

// GetURLPostBody returns a HTTP post body as a byte array from a
// URL, body type and an io.Reader.
func GetURLPostBody(absoluteUrl string, bodyType string, reqBody io.Reader) ([]byte, error) {
	client := &http.Client{}
	res, err := client.Post(absoluteUrl, bodyType, reqBody)
	if err != nil {
		return []byte{}, err
	}
	defer res.Body.Close()
	return ioutil.ReadAll(res.Body)
}

// JoinAbsolute performs a path.Join() while preserving two slashes after the scheme.
func JoinAbsolute(elem ...string) string {
	return regexp.MustCompile(`^([A-Za-z]+:/)`).ReplaceAllString(path.Join(elem...), "${1}/")
}
<|endoftext|>


<commit_before>// Package tollbooth provides rate-limiting logic to HTTP request handler.
package tollbooth

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

	"github.com/didip/tollbooth/errors"
	"github.com/didip/tollbooth/libstring"
	"github.com/didip/tollbooth/limiter"
)

// NewLimiter is a convenience function to limiter.New.
func NewLimiter(max int64, ttl time.Duration, tbOptions *limiter.ExpirableOptions) *limiter.Limiter {
	return limiter.New(tbOptions).SetMax(max).SetTTL(ttl)
}

// LimitByKeys keeps track number of request made by keys separated by pipe.
// It returns HTTPError when limit is exceeded.
func LimitByKeys(limiter *limiter.Limiter, keys []string) *errors.HTTPError {
	if limiter.LimitReached(strings.Join(keys, "|")) {
		return &errors.HTTPError{Message: limiter.GetMessage(), StatusCode: limiter.GetStatusCode()}
	}

	return nil
}

// LimitByRequest builds keys based on http.Request struct,
// loops through all the keys, and check if any one of them returns HTTPError.
func LimitByRequest(limiter *limiter.Limiter, r *http.Request) *errors.HTTPError {
	sliceKeys := BuildKeys(limiter, r)

	// Loop sliceKeys and check if one of them has error.
	for _, keys := range sliceKeys {
		httpError := LimitByKeys(limiter, keys)
		if httpError != nil {
			return httpError
		}
	}

	return nil
}

// BuildKeys generates a slice of keys to rate-limit by given limiter and request structs.
func BuildKeys(limiter *limiter.Limiter, r *http.Request) [][]string {
	remoteIP := libstring.RemoteIP(limiter.GetIPLookups(), limiter.GetForwardedForIndexFromBehind(), r)
	path := r.URL.Path
	sliceKeys := make([][]string, 0)

	// Don't BuildKeys if remoteIP is blank.
	if remoteIP == "" {
		return sliceKeys
	}

	limiterMethods := limiter.GetMethods()
	limiterHeaders := limiter.GetHeaders()
	limiterBasicAuthUsers := limiter.GetBasicAuthUsers()

	limiterHeadersIsSet := limiterHeaders != nil && len(limiterHeaders) > 0
	limiterBasicAuthUsersIsSet := limiterBasicAuthUsers != nil && len(limiterBasicAuthUsers) > 0

	if limiterMethods != nil && limiterHeadersIsSet && limiterBasicAuthUsersIsSet {
		// Limit by HTTP methods and HTTP headers+values and Basic Auth credentials.
		if libstring.StringInSlice(limiterMethods, r.Method) {
			for headerKey, headerValues := range limiterHeaders {
				if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
					// If header values are empty, rate-limit all request with headerKey.
					username, _, ok := r.BasicAuth()
					if ok && libstring.StringInSlice(limiterBasicAuthUsers, username) {
						sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, username})
					}

				} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
					// If header values are not empty, rate-limit all request with headerKey and headerValues.
					for _, headerValue := range headerValues {
						username, _, ok := r.BasicAuth()
						if ok && libstring.StringInSlice(limiterBasicAuthUsers, username) {
							sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, headerValue, username})
						}
					}
				}
			}
		}

	} else if limiterMethods != nil && limiterHeadersIsSet {
		// Limit by HTTP methods and HTTP headers+values.
		if libstring.StringInSlice(limiterMethods, r.Method) {
			for headerKey, headerValues := range limiterHeaders {
				if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
					// If header values are empty, rate-limit all request with headerKey.
					sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey})

				} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
					// If header values are not empty, rate-limit all request with headerKey and headerValues.
					for _, headerValue := range headerValues {
						sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, headerValue})
					}
				}
			}
		}

	} else if limiterMethods != nil && limiterBasicAuthUsersIsSet {
		// Limit by HTTP methods and Basic Auth credentials.
		if libstring.StringInSlice(limiterMethods, r.Method) {
			username, _, ok := r.BasicAuth()
			if ok && libstring.StringInSlice(limiterBasicAuthUsers, username) {
				sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, username})
			}
		}

	} else if limiterMethods != nil {
		// Limit by HTTP methods.
		if libstring.StringInSlice(limiterMethods, r.Method) {
			sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method})
		}

	} else if limiterHeadersIsSet {
		// Limit by HTTP headers+values.
		for headerKey, headerValues := range limiterHeaders {
			if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
				// If header values are empty, rate-limit all request with headerKey.
				sliceKeys = append(sliceKeys, []string{remoteIP, path, headerKey})

			} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
				// If header values are not empty, rate-limit all request with headerKey and headerValues.
				for _, headerValue := range headerValues {
					sliceKeys = append(sliceKeys, []string{remoteIP, path, headerKey, headerValue})
				}
			}
		}

	} else if limiterBasicAuthUsersIsSet {
		// Limit by Basic Auth credentials.
		username, _, ok := r.BasicAuth()
		if ok && libstring.StringInSlice(limiterBasicAuthUsers, username) {
			sliceKeys = append(sliceKeys, []string{remoteIP, path, username})
		}
	} else {
		// Default: Limit by remoteIP and path.
		sliceKeys = append(sliceKeys, []string{remoteIP, path})
	}

	return sliceKeys
}

// setResponseHeaders configures X-Rate-Limit-Limit and X-Rate-Limit-Duration
func setResponseHeaders(lmt *limiter.Limiter, w http.ResponseWriter, r *http.Request) {
	w.Header().Add("X-Rate-Limit-Limit", strconv.FormatInt(lmt.GetMax(), 10))
	w.Header().Add("X-Rate-Limit-Duration", lmt.GetTTL().String())
	w.Header().Add("X-Rate-Limit-Request-Forwarded-For", r.Header.Get("X-Forwarded-For"))
	w.Header().Add("X-Rate-Limit-Request-Remote-Addr", r.RemoteAddr)
}

// LimitHandler is a middleware that performs rate-limiting given http.Handler struct.
func LimitHandler(lmt *limiter.Limiter, next http.Handler) http.Handler {
	middle := func(w http.ResponseWriter, r *http.Request) {
		setResponseHeaders(lmt, w, r)

		httpError := LimitByRequest(lmt, r)
		if httpError != nil {
			w.Header().Add("Content-Type", lmt.GetMessageContentType())
			w.WriteHeader(httpError.StatusCode)
			w.Write([]byte(httpError.Message))

			lmt.ExecOnLimitReached(w, r)
			return
		}

		// There's no rate-limit error, serve the next handler.
		next.ServeHTTP(w, r)
	}

	return http.HandlerFunc(middle)
}

// LimitFuncHandler is a middleware that performs rate-limiting given request handler function.
func LimitFuncHandler(lmt *limiter.Limiter, nextFunc func(http.ResponseWriter, *http.Request)) http.Handler {
	return LimitHandler(lmt, http.HandlerFunc(nextFunc))
}
<commit_msg>LimitBtRequest is not capable of setting response headers.<commit_after>// Package tollbooth provides rate-limiting logic to HTTP request handler.
package tollbooth

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

	"github.com/didip/tollbooth/errors"
	"github.com/didip/tollbooth/libstring"
	"github.com/didip/tollbooth/limiter"
)

// setResponseHeaders configures X-Rate-Limit-Limit and X-Rate-Limit-Duration
func setResponseHeaders(lmt *limiter.Limiter, w http.ResponseWriter, r *http.Request) {
	w.Header().Add("X-Rate-Limit-Limit", strconv.FormatInt(lmt.GetMax(), 10))
	w.Header().Add("X-Rate-Limit-Duration", lmt.GetTTL().String())
	w.Header().Add("X-Rate-Limit-Request-Forwarded-For", r.Header.Get("X-Forwarded-For"))
	w.Header().Add("X-Rate-Limit-Request-Remote-Addr", r.RemoteAddr)
}

// NewLimiter is a convenience function to limiter.New.
func NewLimiter(max int64, ttl time.Duration, tbOptions *limiter.ExpirableOptions) *limiter.Limiter {
	return limiter.New(tbOptions).SetMax(max).SetTTL(ttl)
}

// LimitByKeys keeps track number of request made by keys separated by pipe.
// It returns HTTPError when limit is exceeded.
func LimitByKeys(lmt *limiter.Limiter, keys []string) *errors.HTTPError {
	if lmt.LimitReached(strings.Join(keys, "|")) {
		return &errors.HTTPError{Message: lmt.GetMessage(), StatusCode: lmt.GetStatusCode()}
	}

	return nil
}

// BuildKeys generates a slice of keys to rate-limit by given limiter and request structs.
func BuildKeys(lmt *limiter.Limiter, r *http.Request) [][]string {
	remoteIP := libstring.RemoteIP(lmt.GetIPLookups(), lmt.GetForwardedForIndexFromBehind(), r)
	path := r.URL.Path
	sliceKeys := make([][]string, 0)

	// Don't BuildKeys if remoteIP is blank.
	if remoteIP == "" {
		return sliceKeys
	}

	lmtMethods := lmt.GetMethods()
	lmtHeaders := lmt.GetHeaders()
	lmtBasicAuthUsers := lmt.GetBasicAuthUsers()

	lmtHeadersIsSet := lmtHeaders != nil && len(lmtHeaders) > 0
	lmtBasicAuthUsersIsSet := lmtBasicAuthUsers != nil && len(lmtBasicAuthUsers) > 0

	if lmtMethods != nil && lmtHeadersIsSet && lmtBasicAuthUsersIsSet {
		// Limit by HTTP methods and HTTP headers+values and Basic Auth credentials.
		if libstring.StringInSlice(lmtMethods, r.Method) {
			for headerKey, headerValues := range lmtHeaders {
				if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
					// If header values are empty, rate-limit all request with headerKey.
					username, _, ok := r.BasicAuth()
					if ok && libstring.StringInSlice(lmtBasicAuthUsers, username) {
						sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, username})
					}

				} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
					// If header values are not empty, rate-limit all request with headerKey and headerValues.
					for _, headerValue := range headerValues {
						username, _, ok := r.BasicAuth()
						if ok && libstring.StringInSlice(lmtBasicAuthUsers, username) {
							sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, headerValue, username})
						}
					}
				}
			}
		}

	} else if lmtMethods != nil && lmtHeadersIsSet {
		// Limit by HTTP methods and HTTP headers+values.
		if libstring.StringInSlice(lmtMethods, r.Method) {
			for headerKey, headerValues := range lmtHeaders {
				if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
					// If header values are empty, rate-limit all request with headerKey.
					sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey})

				} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
					// If header values are not empty, rate-limit all request with headerKey and headerValues.
					for _, headerValue := range headerValues {
						sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, headerValue})
					}
				}
			}
		}

	} else if lmtMethods != nil && lmtBasicAuthUsersIsSet {
		// Limit by HTTP methods and Basic Auth credentials.
		if libstring.StringInSlice(lmtMethods, r.Method) {
			username, _, ok := r.BasicAuth()
			if ok && libstring.StringInSlice(lmtBasicAuthUsers, username) {
				sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, username})
			}
		}

	} else if lmtMethods != nil {
		// Limit by HTTP methods.
		if libstring.StringInSlice(lmtMethods, r.Method) {
			sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method})
		}

	} else if lmtHeadersIsSet {
		// Limit by HTTP headers+values.
		for headerKey, headerValues := range lmtHeaders {
			if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
				// If header values are empty, rate-limit all request with headerKey.
				sliceKeys = append(sliceKeys, []string{remoteIP, path, headerKey})

			} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
				// If header values are not empty, rate-limit all request with headerKey and headerValues.
				for _, headerValue := range headerValues {
					sliceKeys = append(sliceKeys, []string{remoteIP, path, headerKey, headerValue})
				}
			}
		}

	} else if lmtBasicAuthUsersIsSet {
		// Limit by Basic Auth credentials.
		username, _, ok := r.BasicAuth()
		if ok && libstring.StringInSlice(lmtBasicAuthUsers, username) {
			sliceKeys = append(sliceKeys, []string{remoteIP, path, username})
		}
	} else {
		// Default: Limit by remoteIP and path.
		sliceKeys = append(sliceKeys, []string{remoteIP, path})
	}

	return sliceKeys
}

// LimitByRequest builds keys based on http.Request struct,
// loops through all the keys, and check if any one of them returns HTTPError.
func LimitByRequest(lmt *limiter.Limiter, w http.ResponseWriter, r *http.Request) *errors.HTTPError {
	setResponseHeaders(lmt, w, r)

	sliceKeys := BuildKeys(lmt, r)

	// Loop sliceKeys and check if one of them has error.
	for _, keys := range sliceKeys {
		httpError := LimitByKeys(lmt, keys)
		if httpError != nil {
			return httpError
		}
	}

	return nil
}

// LimitHandler is a middleware that performs rate-limiting given http.Handler struct.
func LimitHandler(lmt *limiter.Limiter, next http.Handler) http.Handler {
	middle := func(w http.ResponseWriter, r *http.Request) {
		httpError := LimitByRequest(lmt, w, r)
		if httpError != nil {
			w.Header().Add("Content-Type", lmt.GetMessageContentType())
			w.WriteHeader(httpError.StatusCode)
			w.Write([]byte(httpError.Message))

			lmt.ExecOnLimitReached(w, r)
			return
		}

		// There's no rate-limit error, serve the next handler.
		next.ServeHTTP(w, r)
	}

	return http.HandlerFunc(middle)
}

// LimitFuncHandler is a middleware that performs rate-limiting given request handler function.
func LimitFuncHandler(lmt *limiter.Limiter, nextFunc func(http.ResponseWriter, *http.Request)) http.Handler {
	return LimitHandler(lmt, http.HandlerFunc(nextFunc))
}
<|endoftext|>


<commit_before>package views

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net/http"
	"strings"

	"github.com/antihax/evedata/services/vanguard"
)

// Bootstrap Token provides access to public information in EVE Online. This is not used by users.
func init() {
	vanguard.AddAuthRoute("bootstrap", "GET", "/X/bootstrapEveAuth", bootstrapEveSSO)
	vanguard.AddAuthRoute("bootstrap", "GET", "/X/bootstrapEveSSOAnswer", bootstrapEveSSOAnswer)
}

func bootstrapEveSSO(w http.ResponseWriter, r *http.Request) {
	s := vanguard.SessionFromContext(r.Context())
	c := vanguard.GlobalsFromContext(r.Context())

	bootstrapScopes := strings.Split("esi-calendar.respond_calendar_events.v1 esi-calendar.read_calendar_events.v1 esi-location.read_location.v1 esi-location.read_ship_type.v1 esi-mail.organize_mail.v1 esi-mail.read_mail.v1 esi-mail.send_mail.v1 esi-skills.read_skills.v1 esi-skills.read_skillqueue.v1 esi-wallet.read_character_wallet.v1 esi-wallet.read_corporation_wallet.v1 esi-search.search_structures.v1 esi-clones.read_clones.v1 esi-characters.read_contacts.v1 esi-universe.read_structures.v1 esi-bookmarks.read_character_bookmarks.v1 esi-killmails.read_killmails.v1 esi-corporations.read_corporation_membership.v1 esi-assets.read_assets.v1 esi-planets.manage_planets.v1 esi-fleets.read_fleet.v1 esi-fleets.write_fleet.v1 esi-ui.open_window.v1 esi-ui.write_waypoint.v1 esi-characters.write_contacts.v1 esi-fittings.read_fittings.v1 esi-fittings.write_fittings.v1 esi-markets.structure_markets.v1 esi-corporations.read_structures.v1 esi-corporations.write_structures.v1 esi-characters.read_loyalty.v1 esi-characters.read_opportunities.v1 esi-characters.read_chat_channels.v1 esi-characters.read_medals.v1 esi-characters.read_standings.v1 esi-characters.read_agents_research.v1 esi-industry.read_character_jobs.v1 esi-markets.read_character_orders.v1 esi-characters.read_blueprints.v1 esi-characters.read_corporation_roles.v1 esi-location.read_online.v1 esi-contracts.read_character_contracts.v1 esi-clones.read_implants.v1 esi-characters.read_fatigue.v1 esi-killmails.read_corporation_killmails.v1 esi-corporations.track_members.v1 esi-wallet.read_corporation_wallets.v1 esi-characters.read_notifications.v1 esi-corporations.read_divisions.v1 esi-corporations.read_contacts.v1 esi-assets.read_corporation_assets.v1 esi-corporations.read_titles.v1 esi-corporations.read_blueprints.v1 esi-bookmarks.read_corporation_bookmarks.v1 esi-contracts.read_corporation_contracts.v1 esi-corporations.read_standings.v1 esi-corporations.read_starbases.v1 esi-industry.read_corporation_jobs.v1 esi-markets.read_corporation_orders.v1 esi-corporations.read_container_logs.v1 esi-industry.read_character_mining.v1 esi-industry.read_corporation_mining.v1 esi-planets.read_customs_offices.v1 esi-corporations.read_facilities.v1 esi-corporations.read_medals.v1 esi-characters.read_titles.v1 esi-alliances.read_contacts.v1 esi-characters.read_fw_stats.v1 esi-corporations.read_fw_stats.v1 esi-corporations.read_outposts.v1 esi-characterstats.read.v1", " ")

	// Make a code to validate on the return
	b := make([]byte, 16)
	rand.Read(b)
	state := base64.URLEncoding.EncodeToString(b)

	// Save the code to our session store to compare later
	s.Values["TOKENstate"] = state
	err := s.Save(r, w)
	if err != nil {
		httpErr(w, err)
		return
	}

	url := c.TokenAuthenticator.AuthorizeURL(state, true, bootstrapScopes)
	http.Redirect(w, r, url, 302)
	httpErrCode(w, nil, http.StatusMovedPermanently)
}

func bootstrapEveSSOAnswer(w http.ResponseWriter, r *http.Request) {
	c := vanguard.GlobalsFromContext(r.Context())

	code := r.FormValue("code")

	tok, err := c.TokenAuthenticator.TokenExchange(code)
	if err != nil {
		httpErr(w, err)
		return
	}

	tokSrc, err := c.TokenAuthenticator.TokenSource(tok)
	if err != nil {
		httpErr(w, err)
		return
	}

	_, err = c.TokenAuthenticator.Verify(tokSrc)
	if err != nil {
		httpErr(w, err)
		return
	}

	if err != nil {
		httpErr(w, err)
		return
	}

	fmt.Fprintf(w, "%+v\n", tok)
}
<commit_msg>Simplify<commit_after>package views

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net/http"
	"strings"

	"github.com/antihax/evedata/services/vanguard"
)

// Bootstrap Token provides access to public information in EVE Online. This is not used by users.
func init() {
	vanguard.AddAuthRoute("bootstrap", "GET", "/X/bootstrapEveAuth", bootstrapEveSSO)
	vanguard.AddAuthRoute("bootstrap", "GET", "/X/bootstrapEveSSOAnswer", bootstrapEveSSOAnswer)
}

func bootstrapEveSSO(w http.ResponseWriter, r *http.Request) {
	s := vanguard.SessionFromContext(r.Context())
	c := vanguard.GlobalsFromContext(r.Context())

	bootstrapScopes := strings.Split("esi-calendar.respond_calendar_events.v1 esi-calendar.read_calendar_events.v1 esi-location.read_location.v1 esi-location.read_ship_type.v1 esi-mail.organize_mail.v1 esi-mail.read_mail.v1 esi-mail.send_mail.v1 esi-skills.read_skills.v1 esi-skills.read_skillqueue.v1 esi-wallet.read_character_wallet.v1 esi-wallet.read_corporation_wallet.v1 esi-search.search_structures.v1 esi-clones.read_clones.v1 esi-characters.read_contacts.v1 esi-universe.read_structures.v1 esi-bookmarks.read_character_bookmarks.v1 esi-killmails.read_killmails.v1 esi-corporations.read_corporation_membership.v1 esi-assets.read_assets.v1 esi-planets.manage_planets.v1 esi-fleets.read_fleet.v1 esi-fleets.write_fleet.v1 esi-ui.open_window.v1 esi-ui.write_waypoint.v1 esi-characters.write_contacts.v1 esi-fittings.read_fittings.v1 esi-fittings.write_fittings.v1 esi-markets.structure_markets.v1 esi-corporations.read_structures.v1 esi-corporations.write_structures.v1 esi-characters.read_loyalty.v1 esi-characters.read_opportunities.v1 esi-characters.read_chat_channels.v1 esi-characters.read_medals.v1 esi-characters.read_standings.v1 esi-characters.read_agents_research.v1 esi-industry.read_character_jobs.v1 esi-markets.read_character_orders.v1 esi-characters.read_blueprints.v1 esi-characters.read_corporation_roles.v1 esi-location.read_online.v1 esi-contracts.read_character_contracts.v1 esi-clones.read_implants.v1 esi-characters.read_fatigue.v1 esi-killmails.read_corporation_killmails.v1 esi-corporations.track_members.v1 esi-wallet.read_corporation_wallets.v1 esi-characters.read_notifications.v1 esi-corporations.read_divisions.v1 esi-corporations.read_contacts.v1 esi-assets.read_corporation_assets.v1 esi-corporations.read_titles.v1 esi-corporations.read_blueprints.v1 esi-bookmarks.read_corporation_bookmarks.v1 esi-contracts.read_corporation_contracts.v1 esi-corporations.read_standings.v1 esi-corporations.read_starbases.v1 esi-industry.read_corporation_jobs.v1 esi-markets.read_corporation_orders.v1 esi-corporations.read_container_logs.v1 esi-industry.read_character_mining.v1 esi-industry.read_corporation_mining.v1 esi-planets.read_customs_offices.v1 esi-corporations.read_facilities.v1 esi-corporations.read_medals.v1 esi-characters.read_titles.v1 esi-alliances.read_contacts.v1 esi-characters.read_fw_stats.v1 esi-corporations.read_fw_stats.v1 esi-corporations.read_outposts.v1 esi-characterstats.read.v1", " ")

	// Make a code to validate on the return
	b := make([]byte, 16)
	rand.Read(b)
	state := base64.URLEncoding.EncodeToString(b)

	// Save the code to our session store to compare later
	s.Values["TOKENstate"] = state
	if err := s.Save(r, w); err != nil {
		httpErr(w, err)
		return
	}

	url := c.TokenAuthenticator.AuthorizeURL(state, true, bootstrapScopes)
	http.Redirect(w, r, url, 302)
	httpErrCode(w, nil, http.StatusMovedPermanently)
}

func bootstrapEveSSOAnswer(w http.ResponseWriter, r *http.Request) {
	c := vanguard.GlobalsFromContext(r.Context())

	code := r.FormValue("code")

	tok, err := c.TokenAuthenticator.TokenExchange(code)
	if err != nil {
		httpErr(w, err)
		return
	}

	tokSrc, err := c.TokenAuthenticator.TokenSource(tok)
	if err != nil {
		httpErr(w, err)
		return
	}

	_, err = c.TokenAuthenticator.Verify(tokSrc)
	if err != nil {
		httpErr(w, err)
		return
	}

	if err != nil {
		httpErr(w, err)
		return
	}

	fmt.Fprintf(w, "%+v\n", tok)
}
<|endoftext|>


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

package system

import (
	"fmt"
	"os"
	"strings"
)

// Add the provide SSH public key to the core user's list of
// authorized keys
func AuthorizeSSHKeys(user string, keysName string, keys []string) error {
	for name, key := range keys {
		keys[name] = strings.TrimSpace(key)
	}

	// join all keys with newlines, ensuring the resulting string
	// also ends with a newline
	joined := fmt.Sprintf("%s\n", strings.Join(keys, "\n"))

	home, err := UserHome(user)
	if err != nil {
		return err
	}

	if _, err = os.Stat(home + "/.ssh"); err != nil {
		if err = os.MkdirAll(home+"/.ssh", os.FileMode(0755)); err != nil {
			return err
		}
	}

	authorized_file := fmt.Sprintf("%s/.ssh/authorized_keys", home)
	f, err := os.OpenFile(authorized_file, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(joined)

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

package system

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

func diffLines(src, dst []string) []string {
	var tgt []string

	mb := map[string]bool{}

	for _, x := range src {
		mb[x] = true
	}

	for _, x := range dst {
		if _, ok := mb[x]; !ok {
			mb[x] = true
		}
	}

	for k, _ := range mb {
		tgt = append(tgt, k)
	}

	return tgt
}

func readLines(path string) ([]string, error) {
	var lines []string

	file, err := os.Open(path)
	if err != nil {
		return lines, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}
	return lines, scanner.Err()
}

// writeLines writes the lines to the given file.
func writeLines(lines []string, path string) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	w := bufio.NewWriter(file)
	for _, line := range lines {
		fmt.Fprintln(w, line)
	}
	return w.Flush()
}

// Add the provide SSH public key to the core user's list of
// authorized keys
func AuthorizeSSHKeys(user string, keysName string, keys []string) error {
	for name, key := range keys {
		keys[name] = strings.TrimSpace(key)
	}

	// join all keys with newlines, ensuring the resulting string
	// also ends with a newline
	// joined := fmt.Sprintf("%s\n", strings.Join(keys, "\n"))

	home, err := UserHome(user)
	if err != nil {
		return err
	}

	if _, err = os.Stat(home + "/.ssh"); err != nil {
		if err = os.MkdirAll(home+"/.ssh", os.FileMode(0755)); err != nil {
			return err
		}
	}

	authorized_file := fmt.Sprintf("%s/.ssh/authorized_keys", home)
	var newkeys []string
	for _, x := range keys {
		newkeys = append(newkeys, strings.Split(x, "\n")...)
	}
	oldkeys, _ := readLines(authorized_file)

	diffkeys := diffLines(oldkeys, newkeys)
	return writeLines(diffkeys, authorized_file)
}
<|endoftext|>