<commit_before>package provider

import (
	"context"
	"fmt"
	"regexp"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/hashicorp/terraform-plugin-framework/tfsdk"
	"github.com/hashicorp/terraform-plugin-framework/types"
	"github.com/hashicorp/terraform-plugin-go/tftypes"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
	"golang.org/x/crypto/bcrypt"
)

func TestAccResourcePasswordBasic(t *testing.T) {
	resource.UnitTest(t, resource.TestCase{
		PreCheck:                 func() { testAccPreCheck(t) },
		ProtoV6ProviderFactories: testAccProtoV6ProviderFactories(),
		Steps: []resource.TestStep{
			{
				Config: `resource "random_password" "basic" {
  							length = 12
						}`,
				Check: resource.ComposeTestCheckFunc(
					testAccResourceStringCheck("random_password.basic", &customLens{
						customLen: 12,
					}),
				),
			},
			{
				ResourceName: "random_password.basic",
				// Usage of ImportStateIdFunc is required as the value passed to the `terraform import` command needs
				// to be the password itself, as the password resource sets ID to "none" and "result" to the password
				// supplied during import.
				ImportStateIdFunc: func(s *terraform.State) (string, error) {
					id := "random_password.basic"
					rs, ok := s.RootModule().Resources[id]
					if !ok {
						return "", fmt.Errorf("not found: %s", id)
					}
					if rs.Primary.ID == "" {
						return "", fmt.Errorf("no ID is set")
					}

					return rs.Primary.Attributes["result"], nil
				},
				ImportState:             true,
				ImportStateVerify:       true,
				ImportStateVerifyIgnore: []string{"bcrypt_hash", "length", "lower", "number", "special", "upper", "min_lower", "min_numeric", "min_special", "min_upper", "override_special"},
			},
		},
	})
}

func TestAccResourcePasswordOverride(t *testing.T) {
	resource.UnitTest(t, resource.TestCase{
		PreCheck:                 func() { testAccPreCheck(t) },
		ProtoV6ProviderFactories: testAccProtoV6ProviderFactories(),
		Steps: []resource.TestStep{
			{
				Config: `resource "random_password" "override" {
							length = 4
							override_special = "!"
							lower = false
							upper = false
							number = false
						}`,
				Check: resource.ComposeTestCheckFunc(
					testAccResourceStringCheck("random_password.override", &customLens{
						customLen: 4,
					}),
					patternMatch("random_password.override", "!!!!"),
				),
			},
		},
	})
}

func TestAccResourcePasswordMin(t *testing.T) {
	resource.UnitTest(t, resource.TestCase{
		PreCheck:                 func() { testAccPreCheck(t) },
		ProtoV6ProviderFactories: testAccProtoV6ProviderFactories(),
		Steps: []resource.TestStep{
			{
				Config: `resource "random_password" "min" {
							length = 12
							override_special = "!#@"
							min_lower = 2
							min_upper = 3
							min_special = 1
							min_numeric = 4
						}`,
				Check: resource.ComposeTestCheckFunc(
					testAccResourceStringCheck("random_password.min", &customLens{
						customLen: 12,
					}),
					regexMatch("random_password.min", regexp.MustCompile(`([a-z])`), 2),
					regexMatch("random_password.min", regexp.MustCompile(`([A-Z])`), 3),
					regexMatch("random_password.min", regexp.MustCompile(`([0-9])`), 4),
					regexMatch("random_password.min", regexp.MustCompile(`([!#@])`), 1),
				),
			},
		},
	})
}

func TestMigratePasswordStateV0toV1(t *testing.T) {
	raw := tftypes.NewValue(tftypes.Object{}, map[string]tftypes.Value{
		"id":               tftypes.NewValue(tftypes.String, "none"),
		"keepers":          tftypes.NewValue(tftypes.Map{ElementType: tftypes.String}, nil),
		"length":           tftypes.NewValue(tftypes.Number, 16),
		"lower":            tftypes.NewValue(tftypes.Bool, true),
		"min_lower":        tftypes.NewValue(tftypes.Number, 0),
		"min_numeric":      tftypes.NewValue(tftypes.Number, 0),
		"min_special":      tftypes.NewValue(tftypes.Number, 0),
		"min_upper":        tftypes.NewValue(tftypes.Number, 0),
		"number":           tftypes.NewValue(tftypes.Bool, true),
		"override_special": tftypes.NewValue(tftypes.String, "!#$%\u0026*()-_=+[]{}\u003c\u003e:?"),
		"result":           tftypes.NewValue(tftypes.String, "DZy_3*tnonj%Q%Yx"),
		"special":          tftypes.NewValue(tftypes.Bool, true),
		"upper":            tftypes.NewValue(tftypes.Bool, true),
	})

	req := tfsdk.UpgradeResourceStateRequest{
		State: &tfsdk.State{
			Raw:    raw,
			Schema: getPasswordSchemaV0(),
		},
	}

	resp := &tfsdk.UpgradeResourceStateResponse{
		State: tfsdk.State{
			Schema: getPasswordSchemaV1(),
		},
	}

	migratePasswordStateV0toV1(context.Background(), req, resp)

	expected := PasswordModelV1{
		ID:              types.String{Value: "none"},
		Keepers:         types.Map{Null: true, ElemType: types.StringType},
		Length:          types.Int64{Value: 16},
		Special:         types.Bool{Value: true},
		Upper:           types.Bool{Value: true},
		Lower:           types.Bool{Value: true},
		Number:          types.Bool{Value: true},
		MinNumeric:      types.Int64{Value: 0},
		MinUpper:        types.Int64{Value: 0},
		MinLower:        types.Int64{Value: 0},
		MinSpecial:      types.Int64{Value: 0},
		OverrideSpecial: types.String{Value: "!#$%\u0026*()-_=+[]{}\u003c\u003e:?"},
		Result:          types.String{Value: "DZy_3*tnonj%Q%Yx"},
	}

	actual := PasswordModelV1{}
	resp.State.Get(context.Background(), &actual)

	err := bcrypt.CompareHashAndPassword([]byte(actual.BcryptHash.Value), []byte(actual.Result.Value))
	if err != nil {
		t.Error(err)
	}

	// Setting actual.BcryptHash to zero value to allow direct comparison of expected and actual.
	actual.BcryptHash = types.String{}

	if !cmp.Equal(expected, actual) {
		t.Errorf("expected: %+v, got: %+v", expected, actual)
	}
}
<commit_msg>Adding error context (#177)<commit_after>package provider

import (
	"context"
	"fmt"
	"regexp"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/hashicorp/terraform-plugin-framework/tfsdk"
	"github.com/hashicorp/terraform-plugin-framework/types"
	"github.com/hashicorp/terraform-plugin-go/tftypes"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
	"golang.org/x/crypto/bcrypt"
)

func TestAccResourcePasswordBasic(t *testing.T) {
	resource.UnitTest(t, resource.TestCase{
		PreCheck:                 func() { testAccPreCheck(t) },
		ProtoV6ProviderFactories: testAccProtoV6ProviderFactories(),
		Steps: []resource.TestStep{
			{
				Config: `resource "random_password" "basic" {
  							length = 12
						}`,
				Check: resource.ComposeTestCheckFunc(
					testAccResourceStringCheck("random_password.basic", &customLens{
						customLen: 12,
					}),
				),
			},
			{
				ResourceName: "random_password.basic",
				// Usage of ImportStateIdFunc is required as the value passed to the `terraform import` command needs
				// to be the password itself, as the password resource sets ID to "none" and "result" to the password
				// supplied during import.
				ImportStateIdFunc: func(s *terraform.State) (string, error) {
					id := "random_password.basic"
					rs, ok := s.RootModule().Resources[id]
					if !ok {
						return "", fmt.Errorf("not found: %s", id)
					}
					if rs.Primary.ID == "" {
						return "", fmt.Errorf("no ID is set")
					}

					return rs.Primary.Attributes["result"], nil
				},
				ImportState:             true,
				ImportStateVerify:       true,
				ImportStateVerifyIgnore: []string{"bcrypt_hash", "length", "lower", "number", "special", "upper", "min_lower", "min_numeric", "min_special", "min_upper", "override_special"},
			},
		},
	})
}

func TestAccResourcePasswordOverride(t *testing.T) {
	resource.UnitTest(t, resource.TestCase{
		PreCheck:                 func() { testAccPreCheck(t) },
		ProtoV6ProviderFactories: testAccProtoV6ProviderFactories(),
		Steps: []resource.TestStep{
			{
				Config: `resource "random_password" "override" {
							length = 4
							override_special = "!"
							lower = false
							upper = false
							number = false
						}`,
				Check: resource.ComposeTestCheckFunc(
					testAccResourceStringCheck("random_password.override", &customLens{
						customLen: 4,
					}),
					patternMatch("random_password.override", "!!!!"),
				),
			},
		},
	})
}

func TestAccResourcePasswordMin(t *testing.T) {
	resource.UnitTest(t, resource.TestCase{
		PreCheck:                 func() { testAccPreCheck(t) },
		ProtoV6ProviderFactories: testAccProtoV6ProviderFactories(),
		Steps: []resource.TestStep{
			{
				Config: `resource "random_password" "min" {
							length = 12
							override_special = "!#@"
							min_lower = 2
							min_upper = 3
							min_special = 1
							min_numeric = 4
						}`,
				Check: resource.ComposeTestCheckFunc(
					testAccResourceStringCheck("random_password.min", &customLens{
						customLen: 12,
					}),
					regexMatch("random_password.min", regexp.MustCompile(`([a-z])`), 2),
					regexMatch("random_password.min", regexp.MustCompile(`([A-Z])`), 3),
					regexMatch("random_password.min", regexp.MustCompile(`([0-9])`), 4),
					regexMatch("random_password.min", regexp.MustCompile(`([!#@])`), 1),
				),
			},
		},
	})
}

func TestMigratePasswordStateV0toV1(t *testing.T) {
	raw := tftypes.NewValue(tftypes.Object{}, map[string]tftypes.Value{
		"id":               tftypes.NewValue(tftypes.String, "none"),
		"keepers":          tftypes.NewValue(tftypes.Map{ElementType: tftypes.String}, nil),
		"length":           tftypes.NewValue(tftypes.Number, 16),
		"lower":            tftypes.NewValue(tftypes.Bool, true),
		"min_lower":        tftypes.NewValue(tftypes.Number, 0),
		"min_numeric":      tftypes.NewValue(tftypes.Number, 0),
		"min_special":      tftypes.NewValue(tftypes.Number, 0),
		"min_upper":        tftypes.NewValue(tftypes.Number, 0),
		"number":           tftypes.NewValue(tftypes.Bool, true),
		"override_special": tftypes.NewValue(tftypes.String, "!#$%\u0026*()-_=+[]{}\u003c\u003e:?"),
		"result":           tftypes.NewValue(tftypes.String, "DZy_3*tnonj%Q%Yx"),
		"special":          tftypes.NewValue(tftypes.Bool, true),
		"upper":            tftypes.NewValue(tftypes.Bool, true),
	})

	req := tfsdk.UpgradeResourceStateRequest{
		State: &tfsdk.State{
			Raw:    raw,
			Schema: getPasswordSchemaV0(),
		},
	}

	resp := &tfsdk.UpgradeResourceStateResponse{
		State: tfsdk.State{
			Schema: getPasswordSchemaV1(),
		},
	}

	migratePasswordStateV0toV1(context.Background(), req, resp)

	expected := PasswordModelV1{
		ID:              types.String{Value: "none"},
		Keepers:         types.Map{Null: true, ElemType: types.StringType},
		Length:          types.Int64{Value: 16},
		Special:         types.Bool{Value: true},
		Upper:           types.Bool{Value: true},
		Lower:           types.Bool{Value: true},
		Number:          types.Bool{Value: true},
		MinNumeric:      types.Int64{Value: 0},
		MinUpper:        types.Int64{Value: 0},
		MinLower:        types.Int64{Value: 0},
		MinSpecial:      types.Int64{Value: 0},
		OverrideSpecial: types.String{Value: "!#$%\u0026*()-_=+[]{}\u003c\u003e:?"},
		Result:          types.String{Value: "DZy_3*tnonj%Q%Yx"},
	}

	actual := PasswordModelV1{}
	resp.State.Get(context.Background(), &actual)

	err := bcrypt.CompareHashAndPassword([]byte(actual.BcryptHash.Value), []byte(actual.Result.Value))
	if err != nil {
		t.Errorf("unexpected bcrypt comparison error: %s", err)
	}

	// Setting actual.BcryptHash to zero value to allow direct comparison of expected and actual.
	actual.BcryptHash = types.String{}

	if !cmp.Equal(expected, actual) {
		t.Errorf("expected: %+v, got: %+v", expected, actual)
	}
}
<|endoftext|>


<commit_before>package central

import (
	"nli-go/lib/mentalese"
)

type EntityDefinitionsExtracter struct {
	dialogContext *DialogContext
}

func NewEntityDefinitionsExtracter(dialogContext *DialogContext) *EntityDefinitionsExtracter {
	return &EntityDefinitionsExtracter{
		dialogContext: dialogContext,
	}
}

func (e *EntityDefinitionsExtracter) Extract(set mentalese.RelationSet) {
	for _, relation := range set {
		if relation.Predicate == mentalese.PredicateQuant {
			definition := e.removeSelfReferences(relation.Arguments[mentalese.QuantRangeSetIndex].TermValueRelationSet)
			e.AddDefinition(relation.Arguments[mentalese.QuantRangeVariableIndex].TermValue, definition)
		}
		for _, argument := range relation.Arguments {
			if argument.IsRelationSet() {
				e.Extract(argument.TermValueRelationSet)
			}
		}
	}
}

func (e *EntityDefinitionsExtracter) removeSelfReferences(set mentalese.RelationSet) mentalese.RelationSet {
	newSet := mentalese.RelationSet{}

	for _, relation := range set {
		if !e.containsSelfReference(relation) {
			newSet = append(newSet, relation)
		}
	}

	return newSet
}

func (e *EntityDefinitionsExtracter) containsSelfReference(relation mentalese.Relation) bool {

	contains := relation.Predicate == mentalese.PredicateReferenceSlot

	for _, argument := range relation.Arguments {
		if argument.IsRelationSet() {
			for _, child := range argument.TermValueRelationSet {
				contains = contains || e.containsSelfReference(child)
			}
		}
	}

	return contains
}

func (e *EntityDefinitionsExtracter) AddDefinition(variable string, definition mentalese.RelationSet) {
	e.dialogContext.EntityDefinitions.Add(variable, definition)
}
<commit_msg>check rule<commit_after>package central

import (
	"nli-go/lib/mentalese"
)

type EntityDefinitionsExtracter struct {
	dialogContext *DialogContext
}

func NewEntityDefinitionsExtracter(dialogContext *DialogContext) *EntityDefinitionsExtracter {
	return &EntityDefinitionsExtracter{
		dialogContext: dialogContext,
	}
}

func (e *EntityDefinitionsExtracter) Extract(set mentalese.RelationSet) {
	for _, relation := range set {
		if relation.Predicate == mentalese.PredicateQuant {
			definition := e.removeSelfReferences(relation.Arguments[mentalese.QuantRangeSetIndex].TermValueRelationSet)
			e.AddDefinition(relation.Arguments[mentalese.QuantRangeVariableIndex].TermValue, definition)
		}
		for _, argument := range relation.Arguments {
			if argument.IsRelationSet() {
				e.Extract(argument.TermValueRelationSet)
			} else if argument.IsRule() {
				e.Extract(argument.TermValueRule.Pattern)
			}
		}
	}
}

func (e *EntityDefinitionsExtracter) removeSelfReferences(set mentalese.RelationSet) mentalese.RelationSet {
	newSet := mentalese.RelationSet{}

	for _, relation := range set {
		if !e.containsSelfReference(relation) {
			newSet = append(newSet, relation)
		}
	}

	return newSet
}

func (e *EntityDefinitionsExtracter) containsSelfReference(relation mentalese.Relation) bool {

	contains := relation.Predicate == mentalese.PredicateReferenceSlot

	for _, argument := range relation.Arguments {
		if argument.IsRelationSet() {
			for _, child := range argument.TermValueRelationSet {
				contains = contains || e.containsSelfReference(child)
			}
		}
	}

	return contains
}

func (e *EntityDefinitionsExtracter) AddDefinition(variable string, definition mentalese.RelationSet) {
	e.dialogContext.EntityDefinitions.Add(variable, definition)
}
<|endoftext|>


<commit_before>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 mail

import (
	"errors"
	"time"

	"github.com/corestoreio/csfw/utils/log"
	"github.com/go-gomail/gomail"
)

var ErrMailChannelClosed = errors.New("The mail channel has been closed.")

// Daemon represents a daemon which must be created via NewDaemon() function
type Daemon struct {
	msgChan  chan *gomail.Message
	dialer   *gomail.Dialer
	sendFunc gomail.SendFunc
	closed   bool
	// SMTPTimeout closes the connection to the SMTP server if no email was
	// sent in the last default 30 seconds.
	SMTPTimeout time.Duration
}

// Start listens to a channel and sends all incoming messages. Errors will be logged.
// Use code snippet:
//		d := NewDaemon(...)
// 		go func(){
//			if err := d.Worker(); err != nil {
// 				panic(err) // for example
// 			}
// 		}()
//		d.Send(*gomail.Message)
//		d.Stop()
func (dm *Daemon) Worker() error {
	if dm.sendFunc != nil {
		return dm.workerSendFunc()
	}
	return dm.workerDial()
}

func (dm *Daemon) workerSendFunc() error {
	if dm.closed {
		return ErrMailChannelClosed
	}

	for {
		select {
		case m, ok := <-dm.msgChan:
			if !ok {
				dm.closed = true
				return nil
			}

			if err := gomail.Send(dm.sendFunc, m); err != nil {
				log.Error("mail.daemon.Start.Send", "err", err, "message", m)
			}
		}
	}
}

func (dm *Daemon) workerDial() error {
	if dm.closed {
		return ErrMailChannelClosed
	}

	var s gomail.SendCloser
	var err error
	open := false
	for {
		select {
		case m, ok := <-dm.msgChan:
			if !ok {
				dm.closed = true
				return nil
			}
			if !open {
				if s, err = dm.dialer.Dial(); err != nil {
					return log.Error("mail.daemon.Start.Dial", "err", err, "message", m)
				}
				open = true
			}
			if err := gomail.Send(s, m); err != nil {
				log.Error("mail.daemon.Start.Send", "err", err, "message", m)
			}
		// Close the connection to the SMTP server if no email was sent in
		// the last n seconds.
		case <-time.After(dm.SMTPTimeout):
			if open {
				if err := s.Close(); err != nil {
					return log.Error("mail.daemon.Start.Close", "err", err)
				}
				open = false
			}
		}
	}
}

// Stop closes the channel stops the daemon
func (dm *Daemon) Stop() error {
	if dm.closed {
		return ErrMailChannelClosed
	}
	close(dm.msgChan)
	dm.closed = true
	return nil
}

// Send sends a mail
func (dm *Daemon) Send(m *gomail.Message) error {
	if dm.closed {
		return ErrMailChannelClosed
	}
	dm.msgChan <- m
	return nil
}

// SendPlain sends a simple plain text email
func (dm *Daemon) SendPlain(from, to, subject, body string) error {
	return dm.sendMsg(from, to, subject, body, false)
}

// SendPlain sends a simple html email
func (dm *Daemon) SendHtml(from, to, subject, body string) error {
	return dm.sendMsg(from, to, subject, body, true)
}

func (dm *Daemon) sendMsg(from, to, subject, body string, isHtml bool) error {
	if dm.closed {
		return ErrMailChannelClosed
	}
	m := gomail.NewMessage()
	m.SetHeader("From", from)
	m.SetHeader("To", to)
	m.SetHeader("Subject", subject)
	contentType := "text/plain"
	if isHtml {
		contentType = "text/html"
	}
	m.SetBody(contentType, body)
	dm.Send(m)
	return nil
}

// Options applies optional arguments to the daemon
// struct. It returns the last set option. More info about the returned function:
// http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
func (dm *Daemon) Option(opts ...DaemonOption) (previous DaemonOption) {
	for _, o := range opts {
		if o != nil {
			previous = o(dm)
		}
	}
	return previous
}

// DaemonOption can be used as an argument in NewDaemon to configure a daemon.
type DaemonOption func(*Daemon) DaemonOption

// DefaultDialer connects to localhost on port 25.
var DefaultDialer = gomail.NewPlainDialer("localhost", 25, "", "")

// SetMessageChannel sets your custom channel to listen to.
func SetMessageChannel(mailChan chan *gomail.Message) DaemonOption {
	return func(da *Daemon) DaemonOption {
		previous := da.msgChan
		da.msgChan = mailChan
		da.closed = false
		return SetMessageChannel(previous)
	}
}

// SetDialer sets a channel to listen to.
func SetDialer(di *gomail.Dialer) DaemonOption {
	if di == nil {
		di = DefaultDialer
	}
	return func(da *Daemon) DaemonOption {
		previous := da.dialer
		da.dialer = di
		return SetDialer(previous)
	}
}

// SetSendFunc lets you implements your email-sending function for e.g.
// to use any other third party API provider. Setting this option
// will remove the dialer. Your implementation must handle timeouts, etc.
func SetSendFunc(sf gomail.SendFunc) DaemonOption {
	return func(da *Daemon) DaemonOption {
		previous := da.sendFunc
		da.sendFunc = sf
		da.dialer = nil
		return SetSendFunc(previous)
	}
}

// NewDaemon creates a new daemon to send default to localhost:25 and creates
// a default unbuffered channel which can be used via the Send*() function.
func NewDaemon(opts ...DaemonOption) *Daemon {
	d := &Daemon{
		dialer:      DefaultDialer,
		SMTPTimeout: time.Second * 30,
	}
	d.Option(opts...)
	if d.msgChan == nil {
		d.msgChan = make(chan *gomail.Message)
	}
	return d
}
<commit_msg>util/mail: Add config.Reader<commit_after>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 mail

import (
	"errors"
	"time"

	"github.com/corestoreio/csfw/config"
	"github.com/corestoreio/csfw/utils/log"
	"github.com/go-gomail/gomail"
)

const (
	PathSmtpHost = "system/smtp/host"
	PathSmtpPort = "system/smtp/port"
)

var ErrMailChannelClosed = errors.New("The mail channel has been closed.")

// Daemon represents a daemon which must be created via NewDaemon() function
type Daemon struct {
	msgChan  chan *gomail.Message
	dialer   *gomail.Dialer
	sendFunc gomail.SendFunc
	closed   bool
	cr       config.Reader
	// SMTPTimeout closes the connection to the SMTP server if no email was
	// sent in the last default 30 seconds.
	SMTPTimeout time.Duration
}

// Start listens to a channel and sends all incoming messages. Errors will be logged.
// Use code snippet:
//		d := NewDaemon(...)
// 		go func(){
//			if err := d.Worker(); err != nil {
// 				panic(err) // for example
// 			}
// 		}()
//		d.Send(*gomail.Message)
//		d.Stop()
func (dm *Daemon) Worker() error {
	if dm.sendFunc != nil {
		return dm.workerSendFunc()
	}
	return dm.workerDial()
}

func (dm *Daemon) workerSendFunc() error {
	if dm.closed {
		return ErrMailChannelClosed
	}

	for {
		select {
		case m, ok := <-dm.msgChan:
			if !ok {
				dm.closed = true
				return nil
			}

			if err := gomail.Send(dm.sendFunc, m); err != nil {
				log.Error("mail.daemon.Start.Send", "err", err, "message", m)
			}
		}
	}
}

func (dm *Daemon) workerDial() error {
	if dm.closed {
		return ErrMailChannelClosed
	}

	var s gomail.SendCloser
	var err error
	open := false
	for {
		select {
		case m, ok := <-dm.msgChan:
			if !ok {
				dm.closed = true
				return nil
			}
			if !open {
				if s, err = dm.dialer.Dial(); err != nil {
					return log.Error("mail.daemon.Start.Dial", "err", err, "message", m)
				}
				open = true
			}
			if err := gomail.Send(s, m); err != nil {
				log.Error("mail.daemon.Start.Send", "err", err, "message", m)
			}
		// Close the connection to the SMTP server if no email was sent in
		// the last n seconds.
		case <-time.After(dm.SMTPTimeout):
			if open {
				if err := s.Close(); err != nil {
					return log.Error("mail.daemon.Start.Close", "err", err)
				}
				open = false
			}
		}
	}
}

// Stop closes the channel stops the daemon
func (dm *Daemon) Stop() error {
	if dm.closed {
		return ErrMailChannelClosed
	}
	close(dm.msgChan)
	dm.closed = true
	return nil
}

// Send sends a mail
func (dm *Daemon) Send(m *gomail.Message) error {
	if dm.closed {
		return ErrMailChannelClosed
	}
	dm.msgChan <- m
	return nil
}

// SendPlain sends a simple plain text email
func (dm *Daemon) SendPlain(from, to, subject, body string) error {
	return dm.sendMsg(from, to, subject, body, false)
}

// SendPlain sends a simple html email
func (dm *Daemon) SendHtml(from, to, subject, body string) error {
	return dm.sendMsg(from, to, subject, body, true)
}

func (dm *Daemon) sendMsg(from, to, subject, body string, isHtml bool) error {
	if dm.closed {
		return ErrMailChannelClosed
	}
	m := gomail.NewMessage()
	m.SetHeader("From", from)
	m.SetHeader("To", to)
	m.SetHeader("Subject", subject)
	contentType := "text/plain"
	if isHtml {
		contentType = "text/html"
	}
	m.SetBody(contentType, body)
	dm.Send(m)
	return nil
}

// Options applies optional arguments to the daemon
// struct. It returns the last set option. More info about the returned function:
// http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
func (dm *Daemon) Option(opts ...DaemonOption) (previous DaemonOption) {
	for _, o := range opts {
		if o != nil {
			previous = o(dm)
		}
	}
	return previous
}

// DaemonOption can be used as an argument in NewDaemon to configure a daemon.
type DaemonOption func(*Daemon) DaemonOption

// DefaultDialer connects to localhost on port 25.
var DefaultDialer = gomail.NewPlainDialer("localhost", 25, "", "")

// SetMessageChannel sets your custom channel to listen to.
func SetMessageChannel(mailChan chan *gomail.Message) DaemonOption {
	return func(da *Daemon) DaemonOption {
		previous := da.msgChan
		da.msgChan = mailChan
		da.closed = false
		return SetMessageChannel(previous)
	}
}

// SetDialer sets a channel to listen to.
func SetDialer(di *gomail.Dialer) DaemonOption {
	if di == nil {
		di = DefaultDialer
	}
	return func(da *Daemon) DaemonOption {
		previous := da.dialer
		da.dialer = di
		return SetDialer(previous)
	}
}

// SetSendFunc lets you implements your email-sending function for e.g.
// to use any other third party API provider. Setting this option
// will remove the dialer. Your implementation must handle timeouts, etc.
func SetSendFunc(sf gomail.SendFunc) DaemonOption {
	return func(da *Daemon) DaemonOption {
		previous := da.sendFunc
		da.sendFunc = sf
		da.dialer = nil
		return SetSendFunc(previous)
	}
}

// SetStoreConfig sets the config.Reader to the Store.
// Default reader is config.DefaultManager
func SetConfig(cr config.Reader) DaemonOption {
	return func(da *Daemon) DaemonOption {
		previous := da.cr
		da.cr = cr
		return SetConfig(previous)
	}
}

// NewDaemon creates a new daemon to send default to localhost:25 and creates
// a default unbuffered channel which can be used via the Send*() function.
func NewDaemon(opts ...DaemonOption) *Daemon {
	d := &Daemon{
		dialer:      DefaultDialer,
		cr:          config.DefaultManager,
		SMTPTimeout: time.Second * 30,
	}
	d.Option(opts...)
	if d.msgChan == nil {
		d.msgChan = make(chan *gomail.Message)
	}
	return d
}
<|endoftext|>


<commit_before>package log

import (
	"fmt"
	"os"
	"time"
)

// assert interface compliance.
var _ Interface = (*Entry)(nil)

// Entry represents a single log entry.
type Entry struct {
	Logger    *Logger   `json:"-"`
	Fields    Fields    `json:"fields"`
	Level     Level     `json:"level"`
	Timestamp time.Time `json:"timestamp"`
	Message   string    `json:"message"`
	start     time.Time
}

// NewEntry returns a new entry for `log`.
func NewEntry(log *Logger) *Entry {
	return &Entry{
		Logger: log,
		Fields: make(Fields),
	}
}

// clone the entry.
func (e *Entry) clone() *Entry {
	return &Entry{
		Logger: e.Logger,
		Fields: e.Fields,
	}
}

// WithFields returns a new entry with `fields` set.
func (e *Entry) WithFields(fields Fielder) *Entry {
	f := Fields{}

	for k, v := range e.Fields {
		f[k] = v
	}

	for k, v := range fields.Fields() {
		f[k] = v
	}

	return &Entry{Logger: e.Logger, Fields: f}
}

// WithField returns a new entry with the `key` and `value` set.
func (e *Entry) WithField(key string, value interface{}) *Entry {
	return e.WithFields(Fields{key: value})
}

// WithError returns a new entry with the "error" set to `err`.
func (e *Entry) WithError(err error) *Entry {
	return e.WithField("error", err.Error())
}

// Debug level message.
func (e *Entry) Debug(msg string) {
	e.Logger.log(DebugLevel, e, msg)
}

// Info level message.
func (e *Entry) Info(msg string) {
	e.Logger.log(InfoLevel, e, msg)
}

// Warn level message.
func (e *Entry) Warn(msg string) {
	e.Logger.log(WarnLevel, e, msg)
}

// Error level message.
func (e *Entry) Error(msg string) {
	e.Logger.log(ErrorLevel, e, msg)
}

// Fatal level message, followed by an exit.
func (e *Entry) Fatal(msg string) {
	e.Logger.log(FatalLevel, e, msg)
	os.Exit(1)
}

// Debugf level formatted message.
func (e *Entry) Debugf(msg string, v ...interface{}) {
	e.Debug(fmt.Sprintf(msg, v...))
}

// Infof level formatted message.
func (e *Entry) Infof(msg string, v ...interface{}) {
	e.Info(fmt.Sprintf(msg, v...))
}

// Warnf level formatted message.
func (e *Entry) Warnf(msg string, v ...interface{}) {
	e.Warn(fmt.Sprintf(msg, v...))
}

// Errorf level formatted message.
func (e *Entry) Errorf(msg string, v ...interface{}) {
	e.Error(fmt.Sprintf(msg, v...))
}

// Fatalf level formatted message, followed by an exit.
func (e *Entry) Fatalf(msg string, v ...interface{}) {
	e.Fatal(fmt.Sprintf(msg, v...))
}

// Trace returns a new entry with a Stop method to fire off
// a corresponding completion log, useful with defer.
func (e *Entry) Trace(msg string) *Entry {
	e.Info(msg)
	v := e.WithFields(e.Fields)
	v.Message = msg
	v.start = time.Now()
	return v
}

// Stop should be used with Trace, to fire off the completion message. When
// an `err` is passed the "error" field is set, and the log level is error.
func (e *Entry) Stop(err *error) {
	if *err == nil {
		e.WithField("duration", time.Since(e.start)).Info(e.Message)
	} else {
		e.WithField("duration", time.Since(e.start)).WithError(*err).Error(e.Message)
	}
}
<commit_msg>move Entry.clone() lower<commit_after>package log

import (
	"fmt"
	"os"
	"time"
)

// assert interface compliance.
var _ Interface = (*Entry)(nil)

// Entry represents a single log entry.
type Entry struct {
	Logger    *Logger   `json:"-"`
	Fields    Fields    `json:"fields"`
	Level     Level     `json:"level"`
	Timestamp time.Time `json:"timestamp"`
	Message   string    `json:"message"`
	start     time.Time
}

// NewEntry returns a new entry for `log`.
func NewEntry(log *Logger) *Entry {
	return &Entry{
		Logger: log,
		Fields: make(Fields),
	}
}

// WithFields returns a new entry with `fields` set.
func (e *Entry) WithFields(fields Fielder) *Entry {
	f := Fields{}

	for k, v := range e.Fields {
		f[k] = v
	}

	for k, v := range fields.Fields() {
		f[k] = v
	}

	return &Entry{Logger: e.Logger, Fields: f}
}

// WithField returns a new entry with the `key` and `value` set.
func (e *Entry) WithField(key string, value interface{}) *Entry {
	return e.WithFields(Fields{key: value})
}

// WithError returns a new entry with the "error" set to `err`.
func (e *Entry) WithError(err error) *Entry {
	return e.WithField("error", err.Error())
}

// Debug level message.
func (e *Entry) Debug(msg string) {
	e.Logger.log(DebugLevel, e, msg)
}

// Info level message.
func (e *Entry) Info(msg string) {
	e.Logger.log(InfoLevel, e, msg)
}

// Warn level message.
func (e *Entry) Warn(msg string) {
	e.Logger.log(WarnLevel, e, msg)
}

// Error level message.
func (e *Entry) Error(msg string) {
	e.Logger.log(ErrorLevel, e, msg)
}

// Fatal level message, followed by an exit.
func (e *Entry) Fatal(msg string) {
	e.Logger.log(FatalLevel, e, msg)
	os.Exit(1)
}

// Debugf level formatted message.
func (e *Entry) Debugf(msg string, v ...interface{}) {
	e.Debug(fmt.Sprintf(msg, v...))
}

// Infof level formatted message.
func (e *Entry) Infof(msg string, v ...interface{}) {
	e.Info(fmt.Sprintf(msg, v...))
}

// Warnf level formatted message.
func (e *Entry) Warnf(msg string, v ...interface{}) {
	e.Warn(fmt.Sprintf(msg, v...))
}

// Errorf level formatted message.
func (e *Entry) Errorf(msg string, v ...interface{}) {
	e.Error(fmt.Sprintf(msg, v...))
}

// Fatalf level formatted message, followed by an exit.
func (e *Entry) Fatalf(msg string, v ...interface{}) {
	e.Fatal(fmt.Sprintf(msg, v...))
}

// Trace returns a new entry with a Stop method to fire off
// a corresponding completion log, useful with defer.
func (e *Entry) Trace(msg string) *Entry {
	e.Info(msg)
	v := e.WithFields(e.Fields)
	v.Message = msg
	v.start = time.Now()
	return v
}

// Stop should be used with Trace, to fire off the completion message. When
// an `err` is passed the "error" field is set, and the log level is error.
func (e *Entry) Stop(err *error) {
	if *err == nil {
		e.WithField("duration", time.Since(e.start)).Info(e.Message)
	} else {
		e.WithField("duration", time.Since(e.start)).WithError(*err).Error(e.Message)
	}
}

// clone the entry.
func (e *Entry) clone() *Entry {
	return &Entry{
		Logger: e.Logger,
		Fields: e.Fields,
	}
}
<|endoftext|>


<commit_before>package tesTaskEngineWorker

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"syscall"
	"tes/ga4gh"
)

const headerSize = int64(102400)

func readFileHead(path string) []byte {
	f, _ := os.Open(path)
	buffer := make([]byte, headerSize)
	l, _ := f.Read(buffer)
	f.Close()
	return buffer[:l]
}

func FindHostPath(bindings []FSBinding, containerPath string) string {
	for _, binding := range bindings {
		if binding.ContainerPath == containerPath {
			return binding.HostPath
		}
	}
	return ""
}

func FindStdin(bindings []FSBinding, containerPath string) (*os.File, error) {
	stdinPath := FindHostPath(bindings, containerPath)
	if stdinPath != "" {
		return os.Open(stdinPath)
	} else {
		return nil, nil
	}
}

// RunJob runs a job.
func RunJob(job *ga4gh_task_exec.Job, mapper FileMapper) error {
	// Modifies the filemapper's jobID
	mapper.Job(job.JobID)

	// Iterates through job.Task.Resources.Volumes and add the volume to mapper.
	for _, disk := range job.Task.Resources.Volumes {
		mapper.AddVolume(job.JobID, disk.Source, disk.MountPoint)
	}

	// MapInput copies the input.Location into input.Path.
	for _, input := range job.Task.Inputs {
		err := mapper.MapInput(job.JobID, input.Location, input.Path, input.Class)
		if err != nil {
			return err
		}
	}

	// MapOutput finds where to output the results, and adds that
	// to Job. It also sets that output path should be the output
	// location once the job is done.
	for _, output := range job.Task.Outputs {
		err := mapper.MapOutput(job.JobID, output.Location, output.Path, output.Class, output.Create)
		if err != nil {
			return err
		}
	}

	// Loops through Docker Tasks, and label them with i (the index).
	for i, dockerTask := range job.Task.Docker {
		stdin, err := FindStdin(mapper.jobs[job.JobID].Bindings, dockerTask.Stdin)
		if err != nil {
			return fmt.Errorf("Error setting up job stdin: %s", err)
			return err
		}

		// Finds stdout path through mapper.TempFile.
		// Takes stdout from Tool, and outputs into a file.
		stdout, err := mapper.TempFile(job.JobID)
		if err != nil {
			return fmt.Errorf("Error setting up job stdout log: %s", err)
		}
		// Finds stderr path through mapper.TempFile.
		// Takes stderr from Tool, and outputs into a file.
		// `stderr` is a stream where systems error is saved.
		// `err` is Go error.
		stderr, err := mapper.TempFile(job.JobID)
		if err != nil {
			// why two returns? will second return actually return?
			return fmt.Errorf("Error setting up job stderr log: %s", err)
		}
		stdoutPath := stdout.Name()
		stderrPath := stderr.Name()

		// `binds` is a slice of the docker run arguments.
		binds := mapper.GetBindings(job.JobID)

		dcmd := DockerCmd{
			ImageName:       dockerTask.ImageName,
			Cmd:             dockerTask.Cmd,
			Binds:           binds,
			Workdir:         dockerTask.Workdir,
			RemoveContainer: true,
			Stdin:           stdin,
			Stdout:          stdout,
			Stderr:          stderr,
		}
		cmdErr := dcmd.Run()
		exitCode := getExitCode(cmdErr)
		log.Printf("Exit code: %d", exitCode)

		stdout.Close()
		stderr.Close()

		// If `Stderr` is supposed to be added to the volume, copy it.
		if len(dockerTask.Stderr) > 0 {
			hstPath := mapper.HostPath(job.JobID, dockerTask.Stderr)
			if len(hstPath) > 0 {
				copyFileContents(stderrPath, hstPath)
			}

		}
		//If `Stdout` is supposed to be added to the volume, copy it.
		if len(dockerTask.Stdout) > 0 {
			hstPath := mapper.HostPath(job.JobID, dockerTask.Stdout)
			if len(hstPath) > 0 {
				copyFileContents(stdoutPath, hstPath)
			}
		}

		stderrText := readFileHead(stderrPath)
		stdoutText := readFileHead(stdoutPath)

		// Send the scheduler service a job status update
		statusReq := &ga4gh_task_ref.UpdateStatusRequest{
			Id:   job.JobID,
			Step: int64(i),
			Log: &ga4gh_task_exec.JobLog{
				Stdout:   string(stdoutText),
				Stderr:   string(stderrText),
				ExitCode: int32(exitCode),
			},
		}
		// TODO context should be created at the top-level and passed down
		ctx := context.Background()
		sched.UpdateJobStatus(ctx, statusReq)

		if cmdErr != nil {
			return cmdErr
		}
	}

	mapper.FinalizeJob(job.JobID)

	return nil
}

// getExitCode gets the exit status (i.e. exit code) from the result of an executed command.
// The exit code is zero if the command completed without error.
func getExitCode(err error) int {
	if err != nil {
		if exiterr, exitOk := err.(*exec.ExitError); exitOk {
			if status, statusOk := exiterr.Sys().(syscall.WaitStatus); statusOk {
				return status.ExitStatus()
			}
		} else {
			log.Printf("Could not determine exit code. Using default -999")
			return -999
		}
	}
	// The error is nil, the command returned successfully, so exit status is 0.
	return 0
}
<commit_msg>Minor cleanup of engine code<commit_after>package tesTaskEngineWorker

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"syscall"
	"tes/ga4gh"
)

const headerSize = int64(102400)

func readFileHead(path string) []byte {
	f, _ := os.Open(path)
	buffer := make([]byte, headerSize)
	l, _ := f.Read(buffer)
	f.Close()
	return buffer[:l]
}

func FindHostPath(bindings []FSBinding, containerPath string) string {
	for _, binding := range bindings {
		if binding.ContainerPath == containerPath {
			return binding.HostPath
		}
	}
	return ""
}

func FindStdin(bindings []FSBinding, containerPath string) (*os.File, error) {
	stdinPath := FindHostPath(bindings, containerPath)
	if stdinPath != "" {
		return os.Open(stdinPath)
	} else {
		return nil, nil
	}
}

// RunJob runs a job.
func RunJob(job *ga4gh_task_exec.Job, mapper FileMapper) error {
	// Modifies the filemapper's jobID
	mapper.Job(job.JobID)

	// Iterates through job.Task.Resources.Volumes and add the volume to mapper.
	for _, disk := range job.Task.Resources.Volumes {
		mapper.AddVolume(job.JobID, disk.Source, disk.MountPoint)
	}

	// MapInput copies the input.Location into input.Path.
	for _, input := range job.Task.Inputs {
		err := mapper.MapInput(job.JobID, input.Location, input.Path, input.Class)
		if err != nil {
			return err
		}
	}

	// MapOutput finds where to output the results, and adds that
	// to Job. It also sets that output path should be the output
	// location once the job is done.
	for _, output := range job.Task.Outputs {
		err := mapper.MapOutput(job.JobID, output.Location, output.Path, output.Class, output.Create)
		if err != nil {
			return err
		}
	}

	for stepNum, dockerTask := range job.Task.Docker {
		stdin, err := FindStdin(mapper.jobs[job.JobID].Bindings, dockerTask.Stdin)
		if err != nil {
			return fmt.Errorf("Error setting up job stdin: %s", err)
			return err
		}

		// Create file for job stdout
		stdout, err := mapper.TempFile(job.JobID)
		if err != nil {
			return fmt.Errorf("Error setting up job stdout log: %s", err)
		}

		// Create file for job stderr
		stderr, err := mapper.TempFile(job.JobID)
		if err != nil {
			return fmt.Errorf("Error setting up job stderr log: %s", err)
		}
		stdoutPath := stdout.Name()
		stderrPath := stderr.Name()

		// `binds` is a slice of the docker run arguments.
		binds := mapper.GetBindings(job.JobID)

		dcmd := DockerCmd{
			ImageName:       dockerTask.ImageName,
			Cmd:             dockerTask.Cmd,
			Binds:           binds,
			Workdir:         dockerTask.Workdir,
			RemoveContainer: true,
			Stdin:           stdin,
			Stdout:          stdout,
			Stderr:          stderr,
		}
		cmdErr := dcmd.Run()
		exitCode := getExitCode(cmdErr)
		log.Printf("Exit code: %d", exitCode)

		stdout.Close()
		stderr.Close()

		// If `Stderr` is supposed to be added to the volume, copy it.
		if dockerTask.Stderr != "" {
			hstPath := mapper.HostPath(job.JobID, dockerTask.Stderr)
			if len(hstPath) > 0 {
				copyFileContents(stderrPath, hstPath)
			}

		}
		//If `Stdout` is supposed to be added to the volume, copy it.
		if dockerTask.Stdout != "" {
			hstPath := mapper.HostPath(job.JobID, dockerTask.Stdout)
			if len(hstPath) > 0 {
				copyFileContents(stdoutPath, hstPath)
			}
		}

		stderrText := readFileHead(stderrPath)
		stdoutText := readFileHead(stdoutPath)

		// Send the scheduler service a job status update
		statusReq := &ga4gh_task_ref.UpdateStatusRequest{
			Id:   job.JobID,
			Step: int64(stepNum),
			Log: &ga4gh_task_exec.JobLog{
				Stdout:   string(stdoutText),
				Stderr:   string(stderrText),
				ExitCode: int32(exitCode),
			},
		}
		// TODO context should be created at the top-level and passed down
		ctx := context.Background()
		sched.UpdateJobStatus(ctx, statusReq)

		if cmdErr != nil {
			return cmdErr
		}
	}

	mapper.FinalizeJob(job.JobID)

	return nil
}

// getExitCode gets the exit status (i.e. exit code) from the result of an executed command.
// The exit code is zero if the command completed without error.
func getExitCode(err error) int {
	if err != nil {
		if exiterr, exitOk := err.(*exec.ExitError); exitOk {
			if status, statusOk := exiterr.Sys().(syscall.WaitStatus); statusOk {
				return status.ExitStatus()
			}
		} else {
			log.Printf("Could not determine exit code. Using default -999")
			return -999
		}
	}
	// The error is nil, the command returned successfully, so exit status is 0.
	return 0
}
<|endoftext|>


<commit_before>package http

import "net/http"

// HostProvider describes something which can yield hosts for transactions,
// and record a given host's success/failure.
type HostProvider interface {
	Get() (host string, err error)
	Put(host string, success bool)
}

// Proxying implements host proxying logic.
func Proxying(hp HostProvider, next Client) Client {
	return &proxying{
		hp:     hp,
		Client: next,
	}
}

type proxying struct {
	hp HostProvider
	Client
}

func (p proxying) Do(req *http.Request) (*http.Response, error) {
	host, err := p.hp.Get()
	if err != nil {
		return nil, err
	}

	req.Host = host
	resp, err := p.Client.Do(req)

	if err == nil {
		p.hp.Put(host, true)
	} else {
		p.hp.Put(host, false)
	}

	return resp, err
}
<commit_msg>More embedding<commit_after>package http

import "net/http"

// HostProvider describes something which can yield hosts for transactions,
// and record a given host's success/failure.
type HostProvider interface {
	Get() (host string, err error)
	Put(host string, success bool)
}

// Proxying implements host proxying logic.
func Proxying(p HostProvider, next Client) Client {
	return &proxying{
		HostProvider: p,
		Client:       next,
	}
}

type proxying struct {
	HostProvider
	Client
}

func (p proxying) Do(req *http.Request) (*http.Response, error) {
	host, err := p.HostProvider.Get()
	if err != nil {
		return nil, err
	}

	req.Host = host
	resp, err := p.Client.Do(req)

	if err == nil {
		p.HostProvider.Put(host, true)
	} else {
		p.HostProvider.Put(host, false)
	}

	return resp, err
}
<|endoftext|>


<commit_before>package http

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/influxdata/platform"
	pctx "github.com/influxdata/platform/context"
	"github.com/influxdata/platform/kit/errors"
	"github.com/julienschmidt/httprouter"
)

// TelegrafHandler is the handler for the telegraf service
type TelegrafHandler struct {
	*httprouter.Router

	TelegrafService            platform.TelegrafConfigStore
	UserResourceMappingService platform.UserResourceMappingService
}

const (
	telegrafsPath            = "/api/v2/telegrafs"
	telegrafsIDPath          = "/api/v2/telegrafs/:id"
	telegrafsIDMembersPath   = "/api/v2/telegrafs/:id/members"
	telegrafsIDMembersIDPath = "/api/v2/telegrafs/:id/members/:userID"
	telegrafsIDOwnersPath    = "/api/v2/telegrafs/:id/owners"
	telegrafsIDOwnersIDPath  = "/api/v2/telegrafs/:id/owners/:userID"
)

// NewTelegrafHandler returns a new instance of TelegrafHandler.
func NewTelegrafHandler(mappingService platform.UserResourceMappingService) *TelegrafHandler {
	h := &TelegrafHandler{
		Router:                     httprouter.New(),
		UserResourceMappingService: mappingService,
	}
	h.HandlerFunc("POST", telegrafsPath, h.handlePostTelegraf)
	h.HandlerFunc("GET", telegrafsPath, h.handleGetTelegrafs)
	h.HandlerFunc("GET", telegrafsIDPath, h.handleGetTelegraf)
	h.HandlerFunc("DELETE", telegrafsIDPath, h.handleDeleteTelegraf)
	h.HandlerFunc("PUT", telegrafsIDPath, h.handlePutTelegraf)

	h.HandlerFunc("POST", telegrafsIDMembersIDPath, newPostMemberHandler(h.UserResourceMappingService, platform.TelegrafResourceType, platform.Member))
	h.HandlerFunc("GET", telegrafsIDMembersIDPath, newGetMembersHandler(h.UserResourceMappingService, platform.Member))
	h.HandlerFunc("DELETE", telegrafsIDMembersIDPath, newDeleteMemberHandler(h.UserResourceMappingService, platform.Member))

	h.HandlerFunc("POST", telegrafsIDOwnersPath, newPostMemberHandler(h.UserResourceMappingService, platform.TelegrafResourceType, platform.Owner))
	h.HandlerFunc("GET", telegrafsIDOwnersPath, newGetMembersHandler(h.UserResourceMappingService, platform.Owner))
	h.HandlerFunc("DELETE", telegrafsIDOwnersIDPath, newDeleteMemberHandler(h.UserResourceMappingService, platform.Owner))

	return h
}

type telegrafLinks struct {
	Self  string `json:"self"`
	Cells string `json:"cells"`
}

type link struct {
	Self string `json:"self"`
}

type telegrafResponse struct {
	*platform.TelegrafConfig
	Links link `json:"links"`
}

type telegrafResponses struct {
	TelegrafConfigs []telegrafResponse `json:"configurations"`
}

func newTelegrafResponse(tc *platform.TelegrafConfig) telegrafResponse {
	return telegrafResponse{
		TelegrafConfig: tc,
		Links: link{
			Self: fmt.Sprintf("/api/v2/telegrafs/%s", tc.ID.String()),
		},
	}
}

func newTelegrafResponses(tcs []*platform.TelegrafConfig) telegrafResponses {
	resp := telegrafResponses{
		TelegrafConfigs: make([]telegrafResponse, len(tcs)),
	}
	for i, c := range tcs {
		resp.TelegrafConfigs[i] = newTelegrafResponse(c)
	}
	return resp
}

func decodeGetTelegrafRequest(ctx context.Context, r *http.Request) (i platform.ID, err error) {
	params := httprouter.ParamsFromContext(ctx)
	id := params.ByName("id")
	if id == "" {
		return i, errors.InvalidDataf("url missing id")
	}

	if err := i.DecodeFromString(id); err != nil {
		return i, err
	}
	return i, nil
}

func (h *TelegrafHandler) handleGetTelegrafs(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	filter, err := decodeUserResourceMappingFilter(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}
	tcs, _, err := h.TelegrafService.FindTelegrafConfigs(ctx, *filter)
	if err != nil {
		if err == platform.ErrViewNotFound {
			err = errors.New(err.Error(), errors.NotFound)
		}
		EncodeError(ctx, err, w)
		return
	}
	if err := encodeResponse(ctx, w, http.StatusOK, newTelegrafResponses(tcs)); err != nil {
		EncodeError(ctx, err, w)
		return
	}
}

func (h *TelegrafHandler) handleGetTelegraf(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	id, err := decodeGetTelegrafRequest(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}
	tc, err := h.TelegrafService.FindTelegrafConfigByID(ctx, id)
	if err != nil {
		if err == platform.ErrViewNotFound {
			err = errors.New(err.Error(), errors.NotFound)
		}
		EncodeError(ctx, err, w)
		return
	}

	contentType := r.Header.Get("Content-Type")
	switch contentType {
	case "application/octet-stream":
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s.toml\"", strings.Replace(strings.TrimSpace(tc.Name), " ", "_", -1)))
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(tc.TOML()))
	case "application/toml":
		w.Header().Set("Content-Type", "application/toml; charset=utf-8")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(tc.TOML()))
	case "application/json":
		fallthrough
	default:
		if err := encodeResponse(ctx, w, http.StatusOK, newTelegrafResponse(tc)); err != nil {
			EncodeError(ctx, err, w)
			return
		}
	}
}

func decodeUserResourceMappingFilter(ctx context.Context, r *http.Request) (*platform.UserResourceMappingFilter, error) {
	urm := new(platform.UserResourceMappingFilter)
	err := json.NewDecoder(r.Body).Decode(urm)
	return urm, err
}

func decodePostTelegrafRequest(ctx context.Context, r *http.Request) (*platform.TelegrafConfig, error) {
	tc := new(platform.TelegrafConfig)
	err := json.NewDecoder(r.Body).Decode(tc)
	return tc, err
}

func decodePutTelegrafRequest(ctx context.Context, r *http.Request) (*platform.TelegrafConfig, error) {
	tc := new(platform.TelegrafConfig)
	if err := json.NewDecoder(r.Body).Decode(tc); err != nil {
		return nil, err
	}
	params := httprouter.ParamsFromContext(ctx)
	id := params.ByName("id")
	if id == "" {
		return nil, errors.InvalidDataf("url missing id")
	}
	i := new(platform.ID)
	if err := i.DecodeFromString(id); err != nil {
		return nil, err
	}
	tc.ID = *i
	return tc, nil
}

// handlePostTelegraf is the HTTP handler for the POST /api/v2/telegrafs route.
func (h *TelegrafHandler) handlePostTelegraf(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	tc, err := decodePostTelegrafRequest(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}
	now := time.Now()
	auth, err := pctx.GetAuthorizer(ctx)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err := h.TelegrafService.CreateTelegrafConfig(ctx, tc, auth.GetUserID(), now); err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err := encodeResponse(ctx, w, http.StatusCreated, newTelegrafResponse(tc)); err != nil {
		EncodeError(ctx, err, w)
		return
	}
}

// handlePutTelegraf is the HTTP handler for the POST /api/v2/telegrafs route.
func (h *TelegrafHandler) handlePutTelegraf(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	tc, err := decodePutTelegrafRequest(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}
	now := time.Now()
	auth, err := pctx.GetAuthorizer(ctx)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}

	tc, err = h.TelegrafService.UpdateTelegrafConfig(ctx, tc.ID, tc, auth.GetUserID(), now)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err := encodeResponse(ctx, w, http.StatusOK, newTelegrafResponse(tc)); err != nil {
		EncodeError(ctx, err, w)
		return
	}
}

func (h *TelegrafHandler) handleDeleteTelegraf(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	i, err := decodeGetTelegrafRequest(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err = h.TelegrafService.DeleteTelegrafConfig(ctx, i); err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err := encodeResponse(ctx, w, http.StatusNoContent, nil); err != nil {
		EncodeError(ctx, err, w)
		return
	}
}
<commit_msg>check Accept header for mime-type<commit_after>package http

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/influxdata/platform"
	pctx "github.com/influxdata/platform/context"
	"github.com/influxdata/platform/kit/errors"
	"github.com/julienschmidt/httprouter"
)

// TelegrafHandler is the handler for the telegraf service
type TelegrafHandler struct {
	*httprouter.Router

	TelegrafService            platform.TelegrafConfigStore
	UserResourceMappingService platform.UserResourceMappingService
}

const (
	telegrafsPath            = "/api/v2/telegrafs"
	telegrafsIDPath          = "/api/v2/telegrafs/:id"
	telegrafsIDMembersPath   = "/api/v2/telegrafs/:id/members"
	telegrafsIDMembersIDPath = "/api/v2/telegrafs/:id/members/:userID"
	telegrafsIDOwnersPath    = "/api/v2/telegrafs/:id/owners"
	telegrafsIDOwnersIDPath  = "/api/v2/telegrafs/:id/owners/:userID"
)

// NewTelegrafHandler returns a new instance of TelegrafHandler.
func NewTelegrafHandler(mappingService platform.UserResourceMappingService) *TelegrafHandler {
	h := &TelegrafHandler{
		Router:                     httprouter.New(),
		UserResourceMappingService: mappingService,
	}
	h.HandlerFunc("POST", telegrafsPath, h.handlePostTelegraf)
	h.HandlerFunc("GET", telegrafsPath, h.handleGetTelegrafs)
	h.HandlerFunc("GET", telegrafsIDPath, h.handleGetTelegraf)
	h.HandlerFunc("DELETE", telegrafsIDPath, h.handleDeleteTelegraf)
	h.HandlerFunc("PUT", telegrafsIDPath, h.handlePutTelegraf)

	h.HandlerFunc("POST", telegrafsIDMembersIDPath, newPostMemberHandler(h.UserResourceMappingService, platform.TelegrafResourceType, platform.Member))
	h.HandlerFunc("GET", telegrafsIDMembersIDPath, newGetMembersHandler(h.UserResourceMappingService, platform.Member))
	h.HandlerFunc("DELETE", telegrafsIDMembersIDPath, newDeleteMemberHandler(h.UserResourceMappingService, platform.Member))

	h.HandlerFunc("POST", telegrafsIDOwnersPath, newPostMemberHandler(h.UserResourceMappingService, platform.TelegrafResourceType, platform.Owner))
	h.HandlerFunc("GET", telegrafsIDOwnersPath, newGetMembersHandler(h.UserResourceMappingService, platform.Owner))
	h.HandlerFunc("DELETE", telegrafsIDOwnersIDPath, newDeleteMemberHandler(h.UserResourceMappingService, platform.Owner))

	return h
}

type telegrafLinks struct {
	Self  string `json:"self"`
	Cells string `json:"cells"`
}

type link struct {
	Self string `json:"self"`
}

type telegrafResponse struct {
	*platform.TelegrafConfig
	Links link `json:"links"`
}

type telegrafResponses struct {
	TelegrafConfigs []telegrafResponse `json:"configurations"`
}

func newTelegrafResponse(tc *platform.TelegrafConfig) telegrafResponse {
	return telegrafResponse{
		TelegrafConfig: tc,
		Links: link{
			Self: fmt.Sprintf("/api/v2/telegrafs/%s", tc.ID.String()),
		},
	}
}

func newTelegrafResponses(tcs []*platform.TelegrafConfig) telegrafResponses {
	resp := telegrafResponses{
		TelegrafConfigs: make([]telegrafResponse, len(tcs)),
	}
	for i, c := range tcs {
		resp.TelegrafConfigs[i] = newTelegrafResponse(c)
	}
	return resp
}

func decodeGetTelegrafRequest(ctx context.Context, r *http.Request) (i platform.ID, err error) {
	params := httprouter.ParamsFromContext(ctx)
	id := params.ByName("id")
	if id == "" {
		return i, errors.InvalidDataf("url missing id")
	}

	if err := i.DecodeFromString(id); err != nil {
		return i, err
	}
	return i, nil
}

func (h *TelegrafHandler) handleGetTelegrafs(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	filter, err := decodeUserResourceMappingFilter(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}
	tcs, _, err := h.TelegrafService.FindTelegrafConfigs(ctx, *filter)
	if err != nil {
		if err == platform.ErrViewNotFound {
			err = errors.New(err.Error(), errors.NotFound)
		}
		EncodeError(ctx, err, w)
		return
	}
	if err := encodeResponse(ctx, w, http.StatusOK, newTelegrafResponses(tcs)); err != nil {
		EncodeError(ctx, err, w)
		return
	}
}

func (h *TelegrafHandler) handleGetTelegraf(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	id, err := decodeGetTelegrafRequest(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}
	tc, err := h.TelegrafService.FindTelegrafConfigByID(ctx, id)
	if err != nil {
		if err == platform.ErrViewNotFound {
			err = errors.New(err.Error(), errors.NotFound)
		}
		EncodeError(ctx, err, w)
		return
	}

	mimeType := r.Header.Get("Accept")
	switch mimeType {
	case "application/octet-stream":
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s.toml\"", strings.Replace(strings.TrimSpace(tc.Name), " ", "_", -1)))
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(tc.TOML()))
	case "application/json":
		if err := encodeResponse(ctx, w, http.StatusOK, newTelegrafResponse(tc)); err != nil {
			EncodeError(ctx, err, w)
			return
		}
	default:
		w.Header().Set("Content-Type", "application/toml; charset=utf-8")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(tc.TOML()))
	}
}

func decodeUserResourceMappingFilter(ctx context.Context, r *http.Request) (*platform.UserResourceMappingFilter, error) {
	urm := new(platform.UserResourceMappingFilter)
	err := json.NewDecoder(r.Body).Decode(urm)
	return urm, err
}

func decodePostTelegrafRequest(ctx context.Context, r *http.Request) (*platform.TelegrafConfig, error) {
	tc := new(platform.TelegrafConfig)
	err := json.NewDecoder(r.Body).Decode(tc)
	return tc, err
}

func decodePutTelegrafRequest(ctx context.Context, r *http.Request) (*platform.TelegrafConfig, error) {
	tc := new(platform.TelegrafConfig)
	if err := json.NewDecoder(r.Body).Decode(tc); err != nil {
		return nil, err
	}
	params := httprouter.ParamsFromContext(ctx)
	id := params.ByName("id")
	if id == "" {
		return nil, errors.InvalidDataf("url missing id")
	}
	i := new(platform.ID)
	if err := i.DecodeFromString(id); err != nil {
		return nil, err
	}
	tc.ID = *i
	return tc, nil
}

// handlePostTelegraf is the HTTP handler for the POST /api/v2/telegrafs route.
func (h *TelegrafHandler) handlePostTelegraf(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	tc, err := decodePostTelegrafRequest(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}
	now := time.Now()
	auth, err := pctx.GetAuthorizer(ctx)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err := h.TelegrafService.CreateTelegrafConfig(ctx, tc, auth.GetUserID(), now); err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err := encodeResponse(ctx, w, http.StatusCreated, newTelegrafResponse(tc)); err != nil {
		EncodeError(ctx, err, w)
		return
	}
}

// handlePutTelegraf is the HTTP handler for the POST /api/v2/telegrafs route.
func (h *TelegrafHandler) handlePutTelegraf(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	tc, err := decodePutTelegrafRequest(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}
	now := time.Now()
	auth, err := pctx.GetAuthorizer(ctx)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}

	tc, err = h.TelegrafService.UpdateTelegrafConfig(ctx, tc.ID, tc, auth.GetUserID(), now)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err := encodeResponse(ctx, w, http.StatusOK, newTelegrafResponse(tc)); err != nil {
		EncodeError(ctx, err, w)
		return
	}
}

func (h *TelegrafHandler) handleDeleteTelegraf(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	i, err := decodeGetTelegrafRequest(ctx, r)
	if err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err = h.TelegrafService.DeleteTelegrafConfig(ctx, i); err != nil {
		EncodeError(ctx, err, w)
		return
	}

	if err := encodeResponse(ctx, w, http.StatusNoContent, nil); err != nil {
		EncodeError(ctx, err, w)
		return
	}
}
<|endoftext|>


<commit_before>package httpfs

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
)

// Type Mount is used to represent a local directory and have functions
// to facilitate the process of reading/writing files to those directories.
type Mount struct {
	BaseDir string
}

func (m *Mount) realPath(u *url.URL) string {
	return filepath.Join(m.BaseDir, filepath.FromSlash(u.Path))
}

// Return the information from the file pointed by the URL
func (m *Mount) InfoFromURL(u *url.URL) (os.FileInfo, error) {
	return os.Stat(m.realPath(u))
}

// Return a file struct ready for reading
func (m *Mount) OpenReadFile(u *url.URL) (*os.File, error) {
	return os.Open(m.realPath(u))
}

// Return the contents of the directory pointed by the given file
func (m *Mount) ReadDir(u *url.URL) ([]os.FileInfo, error) {
	f, err := os.Open(m.realPath(u))
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return f.Readdir(-1)
}

// Create a new file (and all the directory strucutre) or open the
// existing file for writing
func (m *Mount) CreateOrOpenFileForWrite(u *url.URL) (*os.File, error) {
	rp := m.realPath(u)
	stat, err := os.Stat(rp)
	if os.IsNotExist(err) {
		// new file
		d := filepath.Dir(rp)
		err = os.MkdirAll(d, 0644)
		if err != nil {
			return nil, err
		}

		return os.OpenFile(rp, os.O_CREATE, 0644)
	}
	if err != nil {
		return nil, err
	}
	if stat.IsDir() {
		return nil, fmt.Errorf("%v is a directory", u.Path)
	}
	return os.OpenFile(rp, os.O_RDWR, 0644)
}

// type RawFS allow access to read/write the contents of a file
type RawFS struct {
	mount    *Mount
	metaBase string
}

// Expose the raw files over HTTP
//
// GET: read the contents of the file or directory.
// POST: write the contents of the file.
//
// Directories are exposed using pure HTML and HTML5 microformat
// for extra information
func (r *RawFS) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	info, err := r.mount.InfoFromURL(req.URL)

	if req.Method == "POST" {
		r.serveFile(w, req, info)
		return
	}

	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	if info.IsDir() {
		r.serveDir(w, req, info)
	} else {
		r.serveFile(w, req, info)
	}
}

func (r *RawFS) serveDir(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	childs, err := r.mount.ReadDir(req.URL)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	fmt.Fprintf(w,
		`<!doctype html>

	<title>Listing directory: %v</title>

Listing directory: <data itemprop="name">%v</data>

`,
		info.Name(),
		info.Name())

	for _, child := range childs {
		fmt.Fprintf(w, `<li itemscope itemprop="child"><a itemprop="url" href="%v"><span itemprop="name">%v</span></a> Directory? <span itemprop="dir">%v</span> / <a itemprop="metaurl" href="%v">Stat</a></li>`,
			"./"+child.Name(), child.Name(), child.IsDir(), r.metaForChild(child.Name(), req.URL))
	}

	fmt.Fprintf(w,
		`

`)
}

func (r *RawFS) metaForChild(name string, parent *url.URL) *url.URL {
	u := parent.ResolveReference(&url.URL{})
	u.Path = path.Join(r.metaBase, u.Path, name)
	return u
}

func (r *RawFS) serveFile(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	switch req.Method {
	case "GET":
		r.serveGetFile(w, req, info)
	case "POST", "PUT":
		r.servePostFile(w, req, info)
	default:
		if req.Method != "GET" {
			http.Error(w, "Only GET/POST/PUT at this moment", http.StatusMethodNotAllowed)
			return
		}
	}
}

func (r *RawFS) servePostFile(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	file, err := r.mount.CreateOrOpenFileForWrite(req.URL)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer file.Close()
	_, err = io.Copy(file, req.Body)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
}

func (r *RawFS) serveGetFile(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	rw, err := r.mount.OpenReadFile(req.URL)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer rw.Close()

	http.ServeContent(w, req, info.Name(), info.ModTime(), rw)
}

// Type MetaFS expose the attributes of the FS, things like
// modtime, ownership, etc..
//
// At this momento only GET is supported
type MetaFS struct {
	mount   *Mount
	rawBase string
}

// Expose the MetaFS over HTTP
func (m *MetaFS) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	info, err := m.mount.InfoFromURL(req.URL)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	m.printAsHtml(w, req, info)
}

// Return the raw url representing this file
func (m *MetaFS) rawURLFor(meta *url.URL) *url.URL {
	u := meta.ResolveReference(&url.URL{})
	u.Path = path.Join(m.rawBase, meta.Path)
	return u
}

func (m *MetaFS) printAsHtml(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	rawUrl := m.rawURLFor(req.URL)

	fmt.Fprintf(w,
		`<!doctype html>

	<title>Info about: %v</title>

Name

%v

Directory?

%v

Size

%d

Mod time

%v

Raw url
%v</href>

`,
		req.URL.Path,
		info.Name(),
		info.IsDir(),
		info.Size(),
		info.ModTime(),
		rawUrl,
		info.Name())
}

// type IndexFS is used just to return the links to
// MetaFS and RawFS
type IndexFS struct {
	prefix string
}

// Exposes the IndexFS over HTTP
func (i *IndexFS) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintf(w,
		`<!doctype html>

	<title>Root dir</title>

Meta

meta/

Raw

raw/

`)
}

// type HttpFS is used to group the IndexFS, MetaFS, and the RawFS
type HttpFS struct {
	meta *MetaFS
	raw  *RawFS
	idx  *IndexFS
	mux  *http.ServeMux
}

// Create a new HttpFS with read/write access over the given Mount and
// with the given prefix.
//
// Uses should register HttpFS using StripPrefix since the url is taken as-is
// from the http.Request object without any pre-processing.
//
// The sub paths (/meta/, /raw/) are handled internally and the user don't need
// to worry about them.
func NewHttpFS(m *Mount, prefix string) *HttpFS {
	r := &HttpFS{}
	prefix = path.Clean(prefix)
	r.raw = &RawFS{mount: m, metaBase: path.Join(prefix, "/meta/")}
	r.meta = &MetaFS{mount: m, rawBase: path.Join(prefix, "/raw/")}
	r.idx = &IndexFS{prefix: prefix}
	r.mux = http.NewServeMux()
	r.mux.Handle("/meta/", http.StripPrefix("/meta", r.meta))
	r.mux.Handle("/raw/", http.StripPrefix("/raw", r.raw))
	r.mux.Handle("/", r.idx)
	return r
}

// Expose the HttpFS
func (h *HttpFS) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	log.Printf("[httpfs]-[%v]-%v", req.Method, req.URL)
	h.mux.ServeHTTP(w, req)
}
<commit_msg>make go-lint happy<commit_after>package httpfs

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
)

// Mount is used to represent a local directory and have functions
// to facilitate the process of reading/writing files to those directories.
type Mount struct {
	BaseDir string
}

func (m *Mount) realPath(u *url.URL) string {
	return filepath.Join(m.BaseDir, filepath.FromSlash(u.Path))
}

// InfoFromURL return the information from the file pointed by the URL
func (m *Mount) InfoFromURL(u *url.URL) (os.FileInfo, error) {
	return os.Stat(m.realPath(u))
}

// OpenReadFile return a file struct ready for reading
func (m *Mount) OpenReadFile(u *url.URL) (*os.File, error) {
	return os.Open(m.realPath(u))
}

// ReadDir Return the contents of the directory pointed by the given file
func (m *Mount) ReadDir(u *url.URL) ([]os.FileInfo, error) {
	f, err := os.Open(m.realPath(u))
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return f.Readdir(-1)
}

// CreateOrOpenFileForWrite can reate a new file (and all the directory strucutre)
// or open the existing file for writing
func (m *Mount) CreateOrOpenFileForWrite(u *url.URL) (*os.File, error) {
	rp := m.realPath(u)
	stat, err := os.Stat(rp)
	if os.IsNotExist(err) {
		// new file
		d := filepath.Dir(rp)
		err = os.MkdirAll(d, 0644)
		if err != nil {
			return nil, err
		}

		return os.OpenFile(rp, os.O_CREATE, 0644)
	}
	if err != nil {
		return nil, err
	}
	if stat.IsDir() {
		return nil, fmt.Errorf("%v is a directory", u.Path)
	}
	return os.OpenFile(rp, os.O_RDWR, 0644)
}

// RawFS allow access to read/write the contents of a file
type RawFS struct {
	mount    *Mount
	metaBase string
}

// ServeHTTP expose the raw files over HTTP
//
// GET: read the contents of the file or directory.
// POST: write the contents of the file.
//
// Directories are exposed using pure HTML and HTML5 microformat
// for extra information
func (r *RawFS) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	info, err := r.mount.InfoFromURL(req.URL)

	if req.Method == "POST" {
		r.serveFile(w, req, info)
		return
	}

	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	if info.IsDir() {
		r.serveDir(w, req, info)
	} else {
		r.serveFile(w, req, info)
	}
}

func (r *RawFS) serveDir(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	childs, err := r.mount.ReadDir(req.URL)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	fmt.Fprintf(w,
		`<!doctype html>

	<title>Listing directory: %v</title>

Listing directory: <data itemprop="name">%v</data>

`,
		info.Name(),
		info.Name())

	for _, child := range childs {
		fmt.Fprintf(w, `<li itemscope itemprop="child"><a itemprop="url" href="%v"><span itemprop="name">%v</span></a> Directory? <span itemprop="dir">%v</span> / <a itemprop="metaurl" href="%v">Stat</a></li>`,
			"./"+child.Name(), child.Name(), child.IsDir(), r.metaForChild(child.Name(), req.URL))
	}

	fmt.Fprintf(w,
		`

`)
}

func (r *RawFS) metaForChild(name string, parent *url.URL) *url.URL {
	u := parent.ResolveReference(&url.URL{})
	u.Path = path.Join(r.metaBase, u.Path, name)
	return u
}

func (r *RawFS) serveFile(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	switch req.Method {
	case "GET":
		r.serveGetFile(w, req, info)
	case "POST", "PUT":
		r.servePostFile(w, req, info)
	default:
		if req.Method != "GET" {
			http.Error(w, "Only GET/POST/PUT at this moment", http.StatusMethodNotAllowed)
			return
		}
	}
}

func (r *RawFS) servePostFile(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	file, err := r.mount.CreateOrOpenFileForWrite(req.URL)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer file.Close()
	_, err = io.Copy(file, req.Body)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
}

func (r *RawFS) serveGetFile(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	rw, err := r.mount.OpenReadFile(req.URL)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer rw.Close()

	http.ServeContent(w, req, info.Name(), info.ModTime(), rw)
}

// MetaFS expose the attributes of the FS, things like
// modtime, ownership, etc..
//
// At this momento only GET is supported
type MetaFS struct {
	mount   *Mount
	rawBase string
}

// ServeHTTP expose the MetaFS over HTTP
func (m *MetaFS) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	info, err := m.mount.InfoFromURL(req.URL)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	m.printAsHtml(w, req, info)
}

// Return the raw url representing this file
func (m *MetaFS) rawURLFor(meta *url.URL) *url.URL {
	u := meta.ResolveReference(&url.URL{})
	u.Path = path.Join(m.rawBase, meta.Path)
	return u
}

func (m *MetaFS) printAsHtml(w http.ResponseWriter, req *http.Request, info os.FileInfo) {
	rawUrl := m.rawURLFor(req.URL)

	fmt.Fprintf(w,
		`<!doctype html>

	<title>Info about: %v</title>

Name

%v

Directory?

%v

Size

%d

Mod time

%v

Raw url
%v</href>

`,
		req.URL.Path,
		info.Name(),
		info.IsDir(),
		info.Size(),
		info.ModTime(),
		rawUrl,
		info.Name())
}

// IndexFS is used just to return the links to
// MetaFS and RawFS
type IndexFS struct {
	prefix string
}

// ServeHTTP exposes the IndexFS over HTTP
func (i *IndexFS) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintf(w,
		`<!doctype html>

	<title>Root dir</title>

Meta

meta/

Raw

raw/

`)
}

// HttpFS is used to group the IndexFS, MetaFS, and the RawFS
type HttpFS struct {
	meta *MetaFS
	raw  *RawFS
	idx  *IndexFS
	mux  *http.ServeMux
}

// NewHttpFS create a new HttpFS with read/write access over the given Mount and
// with the given prefix.
//
// Uses should register HttpFS using StripPrefix since the url is taken as-is
// from the http.Request object without any pre-processing.
//
// The sub paths (/meta/, /raw/) are handled internally and the user don't need
// to worry about them.
func NewHttpFS(m *Mount, prefix string) *HttpFS {
	r := &HttpFS{}
	prefix = path.Clean(prefix)
	r.raw = &RawFS{mount: m, metaBase: path.Join(prefix, "/meta/")}
	r.meta = &MetaFS{mount: m, rawBase: path.Join(prefix, "/raw/")}
	r.idx = &IndexFS{prefix: prefix}
	r.mux = http.NewServeMux()
	r.mux.Handle("/meta/", http.StripPrefix("/meta", r.meta))
	r.mux.Handle("/raw/", http.StripPrefix("/raw", r.raw))
	r.mux.Handle("/", r.idx)
	return r
}

// Expose the HttpFS
func (h *HttpFS) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	log.Printf("[httpfs]-[%v]-%v", req.Method, req.URL)
	h.mux.ServeHTTP(w, req)
}
<|endoftext|>


<commit_before>package iptables_test

import (
	"errors"
	"fmt"
	"os/exec"
	"time"

	"code.cloudfoundry.org/guardian/kawasaki/iptables"
	fakes "code.cloudfoundry.org/guardian/kawasaki/iptables/iptablesfakes"
	"code.cloudfoundry.org/guardian/pkg/locksmith"
	"github.com/cloudfoundry/gunk/command_runner/fake_command_runner"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("IPTables controller", func() {
	var (
		netnsName          string
		prefix             string
		iptablesController iptables.IPTables
		fakeLocksmith      *FakeLocksmith
		fakeRunner         *fake_command_runner.FakeCommandRunner
	)

	BeforeEach(func() {
		SetDefaultEventuallyTimeout(3 * time.Second)
		netnsName = fmt.Sprintf("ginkgo-netns-%d", GinkgoParallelNode())
		makeNamespace(netnsName)

		fakeRunner = fake_command_runner.New()
		fakeRunner.WhenRunning(fake_command_runner.CommandSpec{},
			func(cmd *exec.Cmd) error {
				if len(cmd.Args) >= 4 && cmd.Args[3] == "panic" {
					panic("ops")
				}
				return wrapCmdInNs(netnsName, cmd).Run()
			},
		)

		fakeLocksmith = NewFakeLocksmith()

		prefix = fmt.Sprintf("g-%d", GinkgoParallelNode())
		iptablesController = iptables.New("/sbin/iptables", "/sbin/iptables-restore", fakeRunner, fakeLocksmith, prefix)
	})

	AfterEach(func() {
		deleteNamespace(netnsName)
	})

	Describe("CreateChain", func() {
		It("creates the chain", func() {
			Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())

			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-L", "test-chain")), GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
		})

		Context("when the table is nat", func() {
			It("creates the nat chain", func() {
				Expect(iptablesController.CreateChain("nat", "test-chain")).To(Succeed())

				sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", "nat", "-L", "test-chain")), GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())
				Eventually(sess).Should(gexec.Exit(0))
			})
		})

		Context("when the chain already exists", func() {
			BeforeEach(func() {
				Expect(iptablesController.CreateChain("nat", "test-chain")).To(Succeed())
			})

			It("returns an error", func() {
				Expect(iptablesController.CreateChain("nat", "test-chain")).NotTo(Succeed())
			})
		})
	})

	Describe("PrependRule", func() {
		It("prepends the rule", func() {
			fakeTCPRule := new(fakes.FakeRule)
			fakeTCPRule.FlagsReturns([]string{"--protocol", "tcp"})
			fakeUDPRule := new(fakes.FakeRule)
			fakeUDPRule.FlagsReturns([]string{"--protocol", "udp"})

			Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())

			Expect(iptablesController.PrependRule("test-chain", fakeTCPRule)).To(Succeed())
			Expect(iptablesController.PrependRule("test-chain", fakeUDPRule)).To(Succeed())

			buff := gbytes.NewBuffer()
			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-S", "test-chain")), buff, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
			Expect(buff).To(gbytes.Say("-A test-chain -p udp\n-A test-chain -p tcp"))
		})

		It("returns an error when the chain does not exist", func() {
			fakeRule := new(fakes.FakeRule)
			fakeRule.FlagsReturns([]string{})

			Expect(iptablesController.PrependRule("test-chain", fakeRule)).NotTo(Succeed())
		})
	})

	Describe("BulkPrependRules", func() {
		It("appends the rules", func() {
			fakeTCPRule := new(fakes.FakeRule)
			fakeTCPRule.FlagsReturns([]string{"--protocol", "tcp"})
			fakeUDPRule := new(fakes.FakeRule)
			fakeUDPRule.FlagsReturns([]string{"--protocol", "udp"})

			Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())
			Expect(iptablesController.BulkPrependRules("test-chain", []iptables.Rule{
				fakeTCPRule,
				fakeUDPRule,
			})).To(Succeed())

			buff := gbytes.NewBuffer()
			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-S", "test-chain")), buff, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
			Expect(buff).To(gbytes.Say("-A test-chain -p udp\n-A test-chain -p tcp"))
		})

		It("returns an error when the chain does not exist", func() {
			fakeRule := new(fakes.FakeRule)
			fakeRule.FlagsReturns([]string{"--protocol", "tcp"})

			Expect(iptablesController.BulkPrependRules("test-chain", []iptables.Rule{fakeRule})).NotTo(Succeed())
		})

		Context("when there are no rules passed", func() {
			It("does nothing", func() {
				Expect(iptablesController.BulkPrependRules("test-chain", []iptables.Rule{})).To(Succeed())
				Expect(fakeRunner.ExecutedCommands()).To(BeZero())
			})
		})
	})

	Describe("DeleteChain", func() {
		BeforeEach(func() {
			Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())
		})

		It("deletes the chain", func() {
			Expect(iptablesController.DeleteChain("filter", "test-chain")).To(Succeed())

			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-L", "test-chain")), GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(1))
		})

		Context("when the table is nat", func() {
			BeforeEach(func() {
				Expect(iptablesController.CreateChain("nat", "test-chain")).To(Succeed())
			})

			It("deletes the nat chain", func() {
				Expect(iptablesController.DeleteChain("nat", "test-chain")).To(Succeed())

				sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", "nat", "-L", "test-chain")), GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())
				Eventually(sess).Should(gexec.Exit(1))
			})
		})

		Context("when the chain does not exist", func() {
			It("does not return an error", func() {
				Expect(iptablesController.DeleteChain("filter", "test-non-existing-chain")).To(Succeed())
			})
		})
	})

	Describe("FlushChain", func() {
		var table string

		BeforeEach(func() {
			table = "filter"
		})

		JustBeforeEach(func() {
			Expect(iptablesController.CreateChain(table, "test-chain")).To(Succeed())

			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-A", "test-chain", "-j", "ACCEPT")), GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
		})

		It("flushes the chain", func() {
			Expect(iptablesController.FlushChain(table, "test-chain")).To(Succeed())

			buff := gbytes.NewBuffer()
			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-S", "test-chain")), buff, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
			Consistently(buff).ShouldNot(gbytes.Say("-A test-chain"))
		})

		Context("when the table is nat", func() {
			BeforeEach(func() {
				table = "nat"
			})

			It("flushes the nat chain", func() {
				Expect(iptablesController.FlushChain(table, "test-chain")).To(Succeed())

				buff := gbytes.NewBuffer()
				sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-S", "test-chain")), buff, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())
				Eventually(sess).Should(gexec.Exit(0))
				Consistently(buff).ShouldNot(gbytes.Say("-A test-chain"))
			})
		})

		Context("when the chain does not exist", func() {
			It("does not return an error", func() {
				Expect(iptablesController.FlushChain("filter", "test-non-existing-chain")).To(Succeed())
			})
		})
	})

	Describe("DeleteChainReferences", func() {
		var table string

		BeforeEach(func() {
			table = "filter"
		})

		JustBeforeEach(func() {
			Expect(iptablesController.CreateChain(table, "test-chain-1")).To(Succeed())
			Expect(iptablesController.CreateChain(table, "test-chain-2")).To(Succeed())

			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-A", "test-chain-1", "-j", "test-chain-2")), GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
		})

		It("deletes the references", func() {
			Expect(iptablesController.DeleteChainReferences(table, "test-chain-1", "test-chain-2")).To(Succeed())

			Eventually(func() string {
				buff := gbytes.NewBuffer()
				sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-S", "test-chain-1")), buff, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())

				Eventually(sess).Should(gexec.Exit(0))

				return string(buff.Contents())
			}).ShouldNot(ContainSubstring("test-chain-2"))
		})
	})

	Describe("Locking Behaviour", func() {
		Context("when something is holding the lock", func() {
			var fakeUnlocker locksmith.Unlocker
			BeforeEach(func() {
				var err error
				fakeUnlocker, err = fakeLocksmith.Lock("/foo/bar")
				Expect(err).NotTo(HaveOccurred())
			})

			It("blocks on any iptables operations until the lock is freed", func() {
				done := make(chan struct{})
				go func() {
					defer GinkgoRecover()
					Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())
					close(done)
				}()

				Consistently(done).ShouldNot(BeClosed())
				fakeUnlocker.Unlock()
				Eventually(done).Should(BeClosed())
			})
		})

		It("should unlock, ensuring future commands can get the lock", func(done Done) {
			Expect(iptablesController.CreateChain("filter", "test-chain-1")).To(Succeed())
			Expect(iptablesController.CreateChain("filter", "test-chain-2")).To(Succeed())
			close(done)
		}, 2.0)

		It("should lock to correct key", func() {
			Expect(iptablesController.CreateChain("filter", "test-chain-1")).To(Succeed())
			Expect(fakeLocksmith.KeyForLastLock()).To(Equal(iptables.LockKey))
		})

		Context("when locking fails", func() {
			BeforeEach(func() {
				fakeLocksmith.LockReturns(nil, errors.New("failed to lock"))
			})

			It("returns the error", func() {
				Expect(iptablesController.CreateChain("filter", "test-chain")).To(MatchError("failed to lock"))
			})
		})

		Context("when running an iptables command fails", func() {
			It("still unlocks", func(done Done) {
				// this is going to fail, because the chain does not exist
				Expect(iptablesController.PrependRule("non-existent-chain", iptables.SingleFilterRule{})).NotTo(Succeed())
				Expect(iptablesController.CreateChain("filter", "test-chain-2")).To(Succeed())
				close(done)
			}, 2.0)
		})

		Context("when running an iptables command panics", func() {
			It("still unlocks", func(done Done) {
				Expect(func() { iptablesController.PrependRule("panic", iptables.SingleFilterRule{}) }).To(Panic())
				Expect(iptablesController.CreateChain("filter", "test-chain-2")).To(Succeed())
				close(done)
			}, 2.0)
		})

		Context("when unlocking fails", func() {
			BeforeEach(func() {
				fakeLocksmith.UnlockReturns(errors.New("failed to unlock"))
			})

			It("returns the error", func() {
				Expect(iptablesController.CreateChain("filter", "test-chain")).To(MatchError("failed to unlock"))
			})
		})
	})
})

func makeNamespace(nsName string) {
	sess, err := gexec.Start(exec.Command("ip", "netns", "add", nsName), GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(sess).Should(gexec.Exit(0))
}

func deleteNamespace(nsName string) {
	sess, err := gexec.Start(exec.Command("ip", "netns", "delete", nsName), GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(sess).Should(gexec.Exit(0))
}

func wrapCmdInNs(nsName string, cmd *exec.Cmd) *exec.Cmd {
	// We wrap iptables with strace to check whether slowness in #145258087
	// is due to iptables being slow or exiting netns being slow.
	wrappedCmd := exec.Command("strace", "ip", "netns", "exec", nsName)
	wrappedCmd.Args = append(wrappedCmd.Args, cmd.Args...)
	wrappedCmd.Stdin = cmd.Stdin
	wrappedCmd.Stdout = cmd.Stdout
	wrappedCmd.Stderr = cmd.Stderr
	return wrappedCmd
}
<commit_msg>Make strace to print out times<commit_after>package iptables_test

import (
	"errors"
	"fmt"
	"os/exec"
	"time"

	"code.cloudfoundry.org/guardian/kawasaki/iptables"
	fakes "code.cloudfoundry.org/guardian/kawasaki/iptables/iptablesfakes"
	"code.cloudfoundry.org/guardian/pkg/locksmith"
	"github.com/cloudfoundry/gunk/command_runner/fake_command_runner"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("IPTables controller", func() {
	var (
		netnsName          string
		prefix             string
		iptablesController iptables.IPTables
		fakeLocksmith      *FakeLocksmith
		fakeRunner         *fake_command_runner.FakeCommandRunner
	)

	BeforeEach(func() {
		SetDefaultEventuallyTimeout(3 * time.Second)
		netnsName = fmt.Sprintf("ginkgo-netns-%d", GinkgoParallelNode())
		makeNamespace(netnsName)

		fakeRunner = fake_command_runner.New()
		fakeRunner.WhenRunning(fake_command_runner.CommandSpec{},
			func(cmd *exec.Cmd) error {
				if len(cmd.Args) >= 4 && cmd.Args[3] == "panic" {
					panic("ops")
				}
				return wrapCmdInNs(netnsName, cmd).Run()
			},
		)

		fakeLocksmith = NewFakeLocksmith()

		prefix = fmt.Sprintf("g-%d", GinkgoParallelNode())
		iptablesController = iptables.New("/sbin/iptables", "/sbin/iptables-restore", fakeRunner, fakeLocksmith, prefix)
	})

	AfterEach(func() {
		deleteNamespace(netnsName)
	})

	Describe("CreateChain", func() {
		It("creates the chain", func() {
			Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())

			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-L", "test-chain")), GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
		})

		Context("when the table is nat", func() {
			It("creates the nat chain", func() {
				Expect(iptablesController.CreateChain("nat", "test-chain")).To(Succeed())

				sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", "nat", "-L", "test-chain")), GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())
				Eventually(sess).Should(gexec.Exit(0))
			})
		})

		Context("when the chain already exists", func() {
			BeforeEach(func() {
				Expect(iptablesController.CreateChain("nat", "test-chain")).To(Succeed())
			})

			It("returns an error", func() {
				Expect(iptablesController.CreateChain("nat", "test-chain")).NotTo(Succeed())
			})
		})
	})

	Describe("PrependRule", func() {
		It("prepends the rule", func() {
			fakeTCPRule := new(fakes.FakeRule)
			fakeTCPRule.FlagsReturns([]string{"--protocol", "tcp"})
			fakeUDPRule := new(fakes.FakeRule)
			fakeUDPRule.FlagsReturns([]string{"--protocol", "udp"})

			Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())

			Expect(iptablesController.PrependRule("test-chain", fakeTCPRule)).To(Succeed())
			Expect(iptablesController.PrependRule("test-chain", fakeUDPRule)).To(Succeed())

			buff := gbytes.NewBuffer()
			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-S", "test-chain")), buff, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
			Expect(buff).To(gbytes.Say("-A test-chain -p udp\n-A test-chain -p tcp"))
		})

		It("returns an error when the chain does not exist", func() {
			fakeRule := new(fakes.FakeRule)
			fakeRule.FlagsReturns([]string{})

			Expect(iptablesController.PrependRule("test-chain", fakeRule)).NotTo(Succeed())
		})
	})

	Describe("BulkPrependRules", func() {
		It("appends the rules", func() {
			fakeTCPRule := new(fakes.FakeRule)
			fakeTCPRule.FlagsReturns([]string{"--protocol", "tcp"})
			fakeUDPRule := new(fakes.FakeRule)
			fakeUDPRule.FlagsReturns([]string{"--protocol", "udp"})

			Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())
			Expect(iptablesController.BulkPrependRules("test-chain", []iptables.Rule{
				fakeTCPRule,
				fakeUDPRule,
			})).To(Succeed())

			buff := gbytes.NewBuffer()
			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-S", "test-chain")), buff, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
			Expect(buff).To(gbytes.Say("-A test-chain -p udp\n-A test-chain -p tcp"))
		})

		It("returns an error when the chain does not exist", func() {
			fakeRule := new(fakes.FakeRule)
			fakeRule.FlagsReturns([]string{"--protocol", "tcp"})

			Expect(iptablesController.BulkPrependRules("test-chain", []iptables.Rule{fakeRule})).NotTo(Succeed())
		})

		Context("when there are no rules passed", func() {
			It("does nothing", func() {
				Expect(iptablesController.BulkPrependRules("test-chain", []iptables.Rule{})).To(Succeed())
				Expect(fakeRunner.ExecutedCommands()).To(BeZero())
			})
		})
	})

	Describe("DeleteChain", func() {
		BeforeEach(func() {
			Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())
		})

		It("deletes the chain", func() {
			Expect(iptablesController.DeleteChain("filter", "test-chain")).To(Succeed())

			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-L", "test-chain")), GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(1))
		})

		Context("when the table is nat", func() {
			BeforeEach(func() {
				Expect(iptablesController.CreateChain("nat", "test-chain")).To(Succeed())
			})

			It("deletes the nat chain", func() {
				Expect(iptablesController.DeleteChain("nat", "test-chain")).To(Succeed())

				sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", "nat", "-L", "test-chain")), GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())
				Eventually(sess).Should(gexec.Exit(1))
			})
		})

		Context("when the chain does not exist", func() {
			It("does not return an error", func() {
				Expect(iptablesController.DeleteChain("filter", "test-non-existing-chain")).To(Succeed())
			})
		})
	})

	Describe("FlushChain", func() {
		var table string

		BeforeEach(func() {
			table = "filter"
		})

		JustBeforeEach(func() {
			Expect(iptablesController.CreateChain(table, "test-chain")).To(Succeed())

			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-A", "test-chain", "-j", "ACCEPT")), GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
		})

		It("flushes the chain", func() {
			Expect(iptablesController.FlushChain(table, "test-chain")).To(Succeed())

			buff := gbytes.NewBuffer()
			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-S", "test-chain")), buff, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
			Consistently(buff).ShouldNot(gbytes.Say("-A test-chain"))
		})

		Context("when the table is nat", func() {
			BeforeEach(func() {
				table = "nat"
			})

			It("flushes the nat chain", func() {
				Expect(iptablesController.FlushChain(table, "test-chain")).To(Succeed())

				buff := gbytes.NewBuffer()
				sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-S", "test-chain")), buff, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())
				Eventually(sess).Should(gexec.Exit(0))
				Consistently(buff).ShouldNot(gbytes.Say("-A test-chain"))
			})
		})

		Context("when the chain does not exist", func() {
			It("does not return an error", func() {
				Expect(iptablesController.FlushChain("filter", "test-non-existing-chain")).To(Succeed())
			})
		})
	})

	Describe("DeleteChainReferences", func() {
		var table string

		BeforeEach(func() {
			table = "filter"
		})

		JustBeforeEach(func() {
			Expect(iptablesController.CreateChain(table, "test-chain-1")).To(Succeed())
			Expect(iptablesController.CreateChain(table, "test-chain-2")).To(Succeed())

			sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-A", "test-chain-1", "-j", "test-chain-2")), GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())
			Eventually(sess).Should(gexec.Exit(0))
		})

		It("deletes the references", func() {
			Expect(iptablesController.DeleteChainReferences(table, "test-chain-1", "test-chain-2")).To(Succeed())

			Eventually(func() string {
				buff := gbytes.NewBuffer()
				sess, err := gexec.Start(wrapCmdInNs(netnsName, exec.Command("iptables", "-t", table, "-S", "test-chain-1")), buff, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())

				Eventually(sess).Should(gexec.Exit(0))

				return string(buff.Contents())
			}).ShouldNot(ContainSubstring("test-chain-2"))
		})
	})

	Describe("Locking Behaviour", func() {
		Context("when something is holding the lock", func() {
			var fakeUnlocker locksmith.Unlocker
			BeforeEach(func() {
				var err error
				fakeUnlocker, err = fakeLocksmith.Lock("/foo/bar")
				Expect(err).NotTo(HaveOccurred())
			})

			It("blocks on any iptables operations until the lock is freed", func() {
				done := make(chan struct{})
				go func() {
					defer GinkgoRecover()
					Expect(iptablesController.CreateChain("filter", "test-chain")).To(Succeed())
					close(done)
				}()

				Consistently(done).ShouldNot(BeClosed())
				fakeUnlocker.Unlock()
				Eventually(done).Should(BeClosed())
			})
		})

		It("should unlock, ensuring future commands can get the lock", func(done Done) {
			Expect(iptablesController.CreateChain("filter", "test-chain-1")).To(Succeed())
			Expect(iptablesController.CreateChain("filter", "test-chain-2")).To(Succeed())
			close(done)
		}, 2.0)

		It("should lock to correct key", func() {
			Expect(iptablesController.CreateChain("filter", "test-chain-1")).To(Succeed())
			Expect(fakeLocksmith.KeyForLastLock()).To(Equal(iptables.LockKey))
		})

		Context("when locking fails", func() {
			BeforeEach(func() {
				fakeLocksmith.LockReturns(nil, errors.New("failed to lock"))
			})

			It("returns the error", func() {
				Expect(iptablesController.CreateChain("filter", "test-chain")).To(MatchError("failed to lock"))
			})
		})

		Context("when running an iptables command fails", func() {
			It("still unlocks", func(done Done) {
				// this is going to fail, because the chain does not exist
				Expect(iptablesController.PrependRule("non-existent-chain", iptables.SingleFilterRule{})).NotTo(Succeed())
				Expect(iptablesController.CreateChain("filter", "test-chain-2")).To(Succeed())
				close(done)
			}, 2.0)
		})

		Context("when running an iptables command panics", func() {
			It("still unlocks", func(done Done) {
				Expect(func() { iptablesController.PrependRule("panic", iptables.SingleFilterRule{}) }).To(Panic())
				Expect(iptablesController.CreateChain("filter", "test-chain-2")).To(Succeed())
				close(done)
			}, 2.0)
		})

		Context("when unlocking fails", func() {
			BeforeEach(func() {
				fakeLocksmith.UnlockReturns(errors.New("failed to unlock"))
			})

			It("returns the error", func() {
				Expect(iptablesController.CreateChain("filter", "test-chain")).To(MatchError("failed to unlock"))
			})
		})
	})
})

func makeNamespace(nsName string) {
	sess, err := gexec.Start(exec.Command("ip", "netns", "add", nsName), GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(sess).Should(gexec.Exit(0))
}

func deleteNamespace(nsName string) {
	sess, err := gexec.Start(exec.Command("ip", "netns", "delete", nsName), GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(sess).Should(gexec.Exit(0))
}

func wrapCmdInNs(nsName string, cmd *exec.Cmd) *exec.Cmd {
	// We wrap iptables with strace to check whether slowness in #145258087
	// is due to iptables being slow or exiting netns being slow.
	wrappedCmd := exec.Command("strace", "-ttT", "ip", "netns", "exec", nsName)
	wrappedCmd.Args = append(wrappedCmd.Args, cmd.Args...)
	wrappedCmd.Stdin = cmd.Stdin
	wrappedCmd.Stdout = cmd.Stdout
	wrappedCmd.Stderr = cmd.Stderr
	return wrappedCmd
}
<|endoftext|>