<commit_before>package commands

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	mh "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multihash"

	cmds "github.com/jbenet/go-ipfs/commands"
	core "github.com/jbenet/go-ipfs/core"
	dag "github.com/jbenet/go-ipfs/merkledag"
)

// ErrObjectTooLarge is returned when too much data was read from stdin. current limit 512k
var ErrObjectTooLarge = errors.New("input object was too large. limit is 512kbytes")

const inputLimit = 512 * 1024

type Node struct {
	Links []Link
	Data  []byte
}

var ObjectCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Interact with ipfs objects",
		ShortDescription: `
'ipfs object' is a plumbing command used to manipulate DAG objects
directly.`,
		Synopsis: `
ipfs object get <key>             - Get the DAG node named by <key>
ipfs object put <data> <encoding> - Stores input, outputs its key
ipfs object data <key>            - Outputs raw bytes in an object
ipfs object links <key>           - Outputs links pointed to by object
ipfs object stat <key>            - Outputs statistics of object
`,
	},

	Subcommands: map[string]*cmds.Command{
		"data":  objectDataCmd,
		"links": objectLinksCmd,
		"get":   objectGetCmd,
		"put":   objectPutCmd,
		"stat":  objectStatCmd,
	},
}

var objectDataCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Outputs the raw bytes in an IPFS object",
		ShortDescription: `
ipfs data is a plumbing command for retreiving the raw bytes stored in
a DAG node. It outputs to stdout, and <key> is a base58 encoded
multihash.
`,
		LongDescription: `
ipfs data is a plumbing command for retreiving the raw bytes stored in
a DAG node. It outputs to stdout, and <key> is a base58 encoded
multihash.

Note that the "--encoding" option does not affect the output, since the
output is the raw data of the object.
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "Key of the object to retrieve, in base58-encoded multihash format"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		key := req.Arguments()[0]
		return objectData(n, key)
	},
}

var objectLinksCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Outputs the links pointed to by the specified object",
		ShortDescription: `
'ipfs object links' is a plumbing command for retreiving the links from
a DAG node. It outputs to stdout, and <key> is a base58 encoded
multihash.
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "Key of the object to retrieve, in base58-encoded multihash format"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		key := req.Arguments()[0]
		return objectLinks(n, key)
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			object := res.Output().(*Object)
			marshalled := marshalLinks(object.Links)
			return strings.NewReader(marshalled), nil
		},
	},
	Type: Object{},
}

var objectGetCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Get and serialize the DAG node named by <key>",
		ShortDescription: `
'ipfs object get' is a plumbing command for retreiving DAG nodes.
It serializes the DAG node to the format specified by the "--encoding"
flag. It outputs to stdout, and <key> is a base58 encoded multihash.
`,
		LongDescription: `
'ipfs object get' is a plumbing command for retreiving DAG nodes.
It serializes the DAG node to the format specified by the "--encoding"
flag. It outputs to stdout, and <key> is a base58 encoded multihash.

This command outputs data in the following encodings:
  * "protobuf"
  * "json"
  * "xml"
(Specified by the "--encoding" or "-enc" flag)`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "Key of the object to retrieve (in base58-encoded multihash format)"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		key := req.Arguments()[0]

		object, err := objectGet(n, key)
		if err != nil {
			return nil, err
		}

		node := &Node{
			Links: make([]Link, len(object.Links)),
			Data:  object.Data,
		}

		for i, link := range object.Links {
			node.Links[i] = Link{
				Hash: link.Hash.B58String(),
				Name: link.Name,
				Size: link.Size,
			}
		}

		return node, nil
	},
	Type: Node{},
	Marshalers: cmds.MarshalerMap{
		cmds.EncodingType("protobuf"): func(res cmds.Response) (io.Reader, error) {
			node := res.Output().(*Node)
			object, err := deserializeNode(node)
			if err != nil {
				return nil, err
			}

			marshaled, err := object.Marshal()
			if err != nil {
				return nil, err
			}
			return bytes.NewReader(marshaled), nil
		},
	},
}

var objectStatCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Get stats for the DAG node named by <key>",
		ShortDescription: `
'ipfs object stat' is a plumbing command to print DAG node statistics.
<key> is a base58 encoded multihash. It outputs to stdout:

	NumLinks        int number of links in link table
	BlockSize       int size of the raw, encoded data
	LinksSize       int size of the links segment
	DataSize        int size of the data segment
	CumulativeSize  int cumulative size of object and its references
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "Key of the object to retrieve (in base58-encoded multihash format)"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		key := req.Arguments()[0]

		object, err := objectGet(n, key)
		if err != nil {
			return nil, err
		}

		ns, err := object.Stat()
		if err != nil {
			return nil, err
		}

		return ns, nil
	},
	Type: dag.NodeStat{},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			ns := res.Output().(dag.NodeStat)

			var buf bytes.Buffer
			w := func(s string, n int) {
				buf.Write([]byte(fmt.Sprintf("%s: %d\n", s, n)))
			}
			w("NumLinks", ns.NumLinks)
			w("BlockSize", ns.BlockSize)
			w("LinksSize", ns.LinksSize)
			w("DataSize", ns.DataSize)
			w("CumulativeSize", ns.CumulativeSize)

			return &buf, nil
		},
	},
}

var objectPutCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Stores input as a DAG object, outputs its key",
		ShortDescription: `
'ipfs object put' is a plumbing command for storing DAG nodes.
It reads from stdin, and the output is a base58 encoded multihash.
`,
		LongDescription: `
'ipfs object put' is a plumbing command for storing DAG nodes.
It reads from stdin, and the output is a base58 encoded multihash.

Data should be in the format specified by <encoding>.
<encoding> may be one of the following:
	* "protobuf"
	* "json"
`,
	},

	Arguments: []cmds.Argument{
		cmds.FileArg("data", true, false, "Data to be stored as a DAG object"),
		cmds.StringArg("encoding", true, false, "Encoding type of <data>, either \"protobuf\" or \"json\""),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		input, err := req.Files().NextFile()
		if err != nil && err != io.EOF {
			return nil, err
		}

		encoding := req.Arguments()[0]

		output, err := objectPut(n, input, encoding)
		if err != nil {
			errType := cmds.ErrNormal
			if err == ErrUnknownObjectEnc {
				errType = cmds.ErrClient
			}
			return nil, cmds.Error{err.Error(), errType}
		}

		return output, nil
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			object := res.Output().(*Object)
			return strings.NewReader("added " + object.Hash), nil
		},
	},
	Type: Object{},
}

// objectData takes a key string and writes out the raw bytes of that node (if there is one)
func objectData(n *core.IpfsNode, key string) (io.Reader, error) {
	dagnode, err := n.Resolver.ResolvePath(key)
	if err != nil {
		return nil, err
	}

	log.Debugf("objectData: found dagnode %q (# of bytes: %d - # links: %d)", key, len(dagnode.Data), len(dagnode.Links))

	return bytes.NewReader(dagnode.Data), nil
}

// objectLinks takes a key string and lists the links it points to
func objectLinks(n *core.IpfsNode, key string) (*Object, error) {
	dagnode, err := n.Resolver.ResolvePath(key)
	if err != nil {
		return nil, err
	}

	log.Debugf("objectLinks: found dagnode %q (# of bytes: %d - # links: %d)", key, len(dagnode.Data), len(dagnode.Links))

	return getOutput(dagnode)
}

// objectGet takes a key string from args and a format option and serializes the dagnode to that format
func objectGet(n *core.IpfsNode, key string) (*dag.Node, error) {
	dagnode, err := n.Resolver.ResolvePath(key)
	if err != nil {
		return nil, err
	}

	log.Debugf("objectGet: found dagnode %q (# of bytes: %d - # links: %d)", key, len(dagnode.Data), len(dagnode.Links))

	return dagnode, nil
}

// objectPut takes a format option, serializes bytes from stdin and updates the dag with that data
func objectPut(n *core.IpfsNode, input io.Reader, encoding string) (*Object, error) {
	var (
		dagnode *dag.Node
		data    []byte
		err     error
	)

	data, err = ioutil.ReadAll(io.LimitReader(input, inputLimit+10))
	if err != nil {
		return nil, err
	}

	if len(data) >= inputLimit {
		return nil, ErrObjectTooLarge
	}

	switch getObjectEnc(encoding) {
	case objectEncodingJSON:
		node := new(Node)
		err = json.Unmarshal(data, node)
		if err != nil {
			return nil, err
		}

		dagnode, err = deserializeNode(node)
		if err != nil {
			return nil, err
		}

	case objectEncodingProtobuf:
		dagnode, err = dag.Decoded(data)

	default:
		return nil, ErrUnknownObjectEnc
	}

	if err != nil {
		return nil, err
	}

	err = addNode(n, dagnode)
	if err != nil {
		return nil, err
	}

	return getOutput(dagnode)
}

// ErrUnknownObjectEnc is returned if a invalid encoding is supplied
var ErrUnknownObjectEnc = errors.New("unknown object encoding")

type objectEncoding string

const (
	objectEncodingJSON     objectEncoding = "json"
	objectEncodingProtobuf                = "protobuf"
)

func getObjectEnc(o interface{}) objectEncoding {
	v, ok := o.(string)
	if !ok {
		// chosen as default because it's human readable
		log.Warning("option is not a string - falling back to json")
		return objectEncodingJSON
	}

	return objectEncoding(v)
}

func getOutput(dagnode *dag.Node) (*Object, error) {
	key, err := dagnode.Key()
	if err != nil {
		return nil, err
	}

	output := &Object{
		Hash:  key.Pretty(),
		Links: make([]Link, len(dagnode.Links)),
	}

	for i, link := range dagnode.Links {
		output.Links[i] = Link{
			Name: link.Name,
			Hash: link.Hash.B58String(),
			Size: link.Size,
		}
	}

	return output, nil
}

// converts the Node object into a real dag.Node
func deserializeNode(node *Node) (*dag.Node, error) {
	dagnode := new(dag.Node)
	dagnode.Data = node.Data
	dagnode.Links = make([]*dag.Link, len(node.Links))
	for i, link := range node.Links {
		hash, err := mh.FromB58String(link.Hash)
		if err != nil {
			return nil, err
		}
		dagnode.Links[i] = &dag.Link{
			Name: link.Name,
			Size: link.Size,
			Hash: hash,
		}
	}

	return dagnode, nil
}
<commit_msg>ipfs block stat bugfix: cast<commit_after>package commands

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	mh "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multihash"

	cmds "github.com/jbenet/go-ipfs/commands"
	core "github.com/jbenet/go-ipfs/core"
	dag "github.com/jbenet/go-ipfs/merkledag"
)

// ErrObjectTooLarge is returned when too much data was read from stdin. current limit 512k
var ErrObjectTooLarge = errors.New("input object was too large. limit is 512kbytes")

const inputLimit = 512 * 1024

type Node struct {
	Links []Link
	Data  []byte
}

var ObjectCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Interact with ipfs objects",
		ShortDescription: `
'ipfs object' is a plumbing command used to manipulate DAG objects
directly.`,
		Synopsis: `
ipfs object get <key>             - Get the DAG node named by <key>
ipfs object put <data> <encoding> - Stores input, outputs its key
ipfs object data <key>            - Outputs raw bytes in an object
ipfs object links <key>           - Outputs links pointed to by object
ipfs object stat <key>            - Outputs statistics of object
`,
	},

	Subcommands: map[string]*cmds.Command{
		"data":  objectDataCmd,
		"links": objectLinksCmd,
		"get":   objectGetCmd,
		"put":   objectPutCmd,
		"stat":  objectStatCmd,
	},
}

var objectDataCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Outputs the raw bytes in an IPFS object",
		ShortDescription: `
ipfs data is a plumbing command for retreiving the raw bytes stored in
a DAG node. It outputs to stdout, and <key> is a base58 encoded
multihash.
`,
		LongDescription: `
ipfs data is a plumbing command for retreiving the raw bytes stored in
a DAG node. It outputs to stdout, and <key> is a base58 encoded
multihash.

Note that the "--encoding" option does not affect the output, since the
output is the raw data of the object.
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "Key of the object to retrieve, in base58-encoded multihash format"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		key := req.Arguments()[0]
		return objectData(n, key)
	},
}

var objectLinksCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Outputs the links pointed to by the specified object",
		ShortDescription: `
'ipfs object links' is a plumbing command for retreiving the links from
a DAG node. It outputs to stdout, and <key> is a base58 encoded
multihash.
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "Key of the object to retrieve, in base58-encoded multihash format"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		key := req.Arguments()[0]
		return objectLinks(n, key)
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			object := res.Output().(*Object)
			marshalled := marshalLinks(object.Links)
			return strings.NewReader(marshalled), nil
		},
	},
	Type: Object{},
}

var objectGetCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Get and serialize the DAG node named by <key>",
		ShortDescription: `
'ipfs object get' is a plumbing command for retreiving DAG nodes.
It serializes the DAG node to the format specified by the "--encoding"
flag. It outputs to stdout, and <key> is a base58 encoded multihash.
`,
		LongDescription: `
'ipfs object get' is a plumbing command for retreiving DAG nodes.
It serializes the DAG node to the format specified by the "--encoding"
flag. It outputs to stdout, and <key> is a base58 encoded multihash.

This command outputs data in the following encodings:
  * "protobuf"
  * "json"
  * "xml"
(Specified by the "--encoding" or "-enc" flag)`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "Key of the object to retrieve (in base58-encoded multihash format)"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		key := req.Arguments()[0]

		object, err := objectGet(n, key)
		if err != nil {
			return nil, err
		}

		node := &Node{
			Links: make([]Link, len(object.Links)),
			Data:  object.Data,
		}

		for i, link := range object.Links {
			node.Links[i] = Link{
				Hash: link.Hash.B58String(),
				Name: link.Name,
				Size: link.Size,
			}
		}

		return node, nil
	},
	Type: Node{},
	Marshalers: cmds.MarshalerMap{
		cmds.EncodingType("protobuf"): func(res cmds.Response) (io.Reader, error) {
			node := res.Output().(*Node)
			object, err := deserializeNode(node)
			if err != nil {
				return nil, err
			}

			marshaled, err := object.Marshal()
			if err != nil {
				return nil, err
			}
			return bytes.NewReader(marshaled), nil
		},
	},
}

var objectStatCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Get stats for the DAG node named by <key>",
		ShortDescription: `
'ipfs object stat' is a plumbing command to print DAG node statistics.
<key> is a base58 encoded multihash. It outputs to stdout:

	NumLinks        int number of links in link table
	BlockSize       int size of the raw, encoded data
	LinksSize       int size of the links segment
	DataSize        int size of the data segment
	CumulativeSize  int cumulative size of object and its references
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "Key of the object to retrieve (in base58-encoded multihash format)"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		key := req.Arguments()[0]

		object, err := objectGet(n, key)
		if err != nil {
			return nil, err
		}

		ns, err := object.Stat()
		if err != nil {
			return nil, err
		}

		return ns, nil
	},
	Type: dag.NodeStat{},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			ns := res.Output().(*dag.NodeStat)

			var buf bytes.Buffer
			w := func(s string, n int) {
				buf.Write([]byte(fmt.Sprintf("%s: %d\n", s, n)))
			}
			w("NumLinks", ns.NumLinks)
			w("BlockSize", ns.BlockSize)
			w("LinksSize", ns.LinksSize)
			w("DataSize", ns.DataSize)
			w("CumulativeSize", ns.CumulativeSize)

			return &buf, nil
		},
	},
}

var objectPutCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Stores input as a DAG object, outputs its key",
		ShortDescription: `
'ipfs object put' is a plumbing command for storing DAG nodes.
It reads from stdin, and the output is a base58 encoded multihash.
`,
		LongDescription: `
'ipfs object put' is a plumbing command for storing DAG nodes.
It reads from stdin, and the output is a base58 encoded multihash.

Data should be in the format specified by <encoding>.
<encoding> may be one of the following:
	* "protobuf"
	* "json"
`,
	},

	Arguments: []cmds.Argument{
		cmds.FileArg("data", true, false, "Data to be stored as a DAG object"),
		cmds.StringArg("encoding", true, false, "Encoding type of <data>, either \"protobuf\" or \"json\""),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		n, err := req.Context().GetNode()
		if err != nil {
			return nil, err
		}

		input, err := req.Files().NextFile()
		if err != nil && err != io.EOF {
			return nil, err
		}

		encoding := req.Arguments()[0]

		output, err := objectPut(n, input, encoding)
		if err != nil {
			errType := cmds.ErrNormal
			if err == ErrUnknownObjectEnc {
				errType = cmds.ErrClient
			}
			return nil, cmds.Error{err.Error(), errType}
		}

		return output, nil
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			object := res.Output().(*Object)
			return strings.NewReader("added " + object.Hash), nil
		},
	},
	Type: Object{},
}

// objectData takes a key string and writes out the raw bytes of that node (if there is one)
func objectData(n *core.IpfsNode, key string) (io.Reader, error) {
	dagnode, err := n.Resolver.ResolvePath(key)
	if err != nil {
		return nil, err
	}

	log.Debugf("objectData: found dagnode %q (# of bytes: %d - # links: %d)", key, len(dagnode.Data), len(dagnode.Links))

	return bytes.NewReader(dagnode.Data), nil
}

// objectLinks takes a key string and lists the links it points to
func objectLinks(n *core.IpfsNode, key string) (*Object, error) {
	dagnode, err := n.Resolver.ResolvePath(key)
	if err != nil {
		return nil, err
	}

	log.Debugf("objectLinks: found dagnode %q (# of bytes: %d - # links: %d)", key, len(dagnode.Data), len(dagnode.Links))

	return getOutput(dagnode)
}

// objectGet takes a key string from args and a format option and serializes the dagnode to that format
func objectGet(n *core.IpfsNode, key string) (*dag.Node, error) {
	dagnode, err := n.Resolver.ResolvePath(key)
	if err != nil {
		return nil, err
	}

	log.Debugf("objectGet: found dagnode %q (# of bytes: %d - # links: %d)", key, len(dagnode.Data), len(dagnode.Links))

	return dagnode, nil
}

// objectPut takes a format option, serializes bytes from stdin and updates the dag with that data
func objectPut(n *core.IpfsNode, input io.Reader, encoding string) (*Object, error) {
	var (
		dagnode *dag.Node
		data    []byte
		err     error
	)

	data, err = ioutil.ReadAll(io.LimitReader(input, inputLimit+10))
	if err != nil {
		return nil, err
	}

	if len(data) >= inputLimit {
		return nil, ErrObjectTooLarge
	}

	switch getObjectEnc(encoding) {
	case objectEncodingJSON:
		node := new(Node)
		err = json.Unmarshal(data, node)
		if err != nil {
			return nil, err
		}

		dagnode, err = deserializeNode(node)
		if err != nil {
			return nil, err
		}

	case objectEncodingProtobuf:
		dagnode, err = dag.Decoded(data)

	default:
		return nil, ErrUnknownObjectEnc
	}

	if err != nil {
		return nil, err
	}

	err = addNode(n, dagnode)
	if err != nil {
		return nil, err
	}

	return getOutput(dagnode)
}

// ErrUnknownObjectEnc is returned if a invalid encoding is supplied
var ErrUnknownObjectEnc = errors.New("unknown object encoding")

type objectEncoding string

const (
	objectEncodingJSON     objectEncoding = "json"
	objectEncodingProtobuf                = "protobuf"
)

func getObjectEnc(o interface{}) objectEncoding {
	v, ok := o.(string)
	if !ok {
		// chosen as default because it's human readable
		log.Warning("option is not a string - falling back to json")
		return objectEncodingJSON
	}

	return objectEncoding(v)
}

func getOutput(dagnode *dag.Node) (*Object, error) {
	key, err := dagnode.Key()
	if err != nil {
		return nil, err
	}

	output := &Object{
		Hash:  key.Pretty(),
		Links: make([]Link, len(dagnode.Links)),
	}

	for i, link := range dagnode.Links {
		output.Links[i] = Link{
			Name: link.Name,
			Hash: link.Hash.B58String(),
			Size: link.Size,
		}
	}

	return output, nil
}

// converts the Node object into a real dag.Node
func deserializeNode(node *Node) (*dag.Node, error) {
	dagnode := new(dag.Node)
	dagnode.Data = node.Data
	dagnode.Links = make([]*dag.Link, len(node.Links))
	for i, link := range node.Links {
		hash, err := mh.FromB58String(link.Hash)
		if err != nil {
			return nil, err
		}
		dagnode.Links[i] = &dag.Link{
			Name: link.Name,
			Size: link.Size,
			Hash: hash,
		}
	}

	return dagnode, nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"io/ioutil"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/mitchellh/go-homedir"

	"errors"

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

func resourceAwsLambdaFunction() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLambdaFunctionCreate,
		Read:   resourceAwsLambdaFunctionRead,
		Update: resourceAwsLambdaFunctionUpdate,
		Delete: resourceAwsLambdaFunctionDelete,

		Schema: map[string]*schema.Schema{
			"filename": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"s3_bucket", "s3_key", "s3_object_version"},
			},
			"s3_bucket": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"s3_key": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"s3_object_version": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"function_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"handler": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"memory_size": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  128,
			},
			"role": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"runtime": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  "nodejs",
			},
			"timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  3,
			},
			"vpc_config": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"subnet_ids": &schema.Schema{
							Type:     schema.TypeSet,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set:      schema.HashString,
						},
						"security_group_ids": &schema.Schema{
							Type:     schema.TypeSet,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set:      schema.HashString,
						},
					},
				},
			},
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"last_modified": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"source_code_hash": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
		},
	}
}

// resourceAwsLambdaFunction maps to:
// CreateFunction in the API / SDK
func resourceAwsLambdaFunctionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	functionName := d.Get("function_name").(string)
	iamRole := d.Get("role").(string)

	log.Printf("[DEBUG] Creating Lambda Function %s with role %s", functionName, iamRole)

	var functionCode *lambda.FunctionCode
	if v, ok := d.GetOk("filename"); ok {
		file, err := loadFileContent(v.(string))
		if err != nil {
			return fmt.Errorf("Unable to load %q: %s", v.(string), err)
		}
		functionCode = &lambda.FunctionCode{
			ZipFile: file,
		}
	} else {
		s3Bucket, bucketOk := d.GetOk("s3_bucket")
		s3Key, keyOk := d.GetOk("s3_key")
		s3ObjectVersion, versionOk := d.GetOk("s3_object_version")
		if !bucketOk || !keyOk {
			return errors.New("s3_bucket and s3_key must all be set while using S3 code source")
		}
		functionCode = &lambda.FunctionCode{
			S3Bucket: aws.String(s3Bucket.(string)),
			S3Key:    aws.String(s3Key.(string)),
		}
		if versionOk {
			functionCode.S3ObjectVersion = aws.String(s3ObjectVersion.(string))
		}
	}

	params := &lambda.CreateFunctionInput{
		Code:         functionCode,
		Description:  aws.String(d.Get("description").(string)),
		FunctionName: aws.String(functionName),
		Handler:      aws.String(d.Get("handler").(string)),
		MemorySize:   aws.Int64(int64(d.Get("memory_size").(int))),
		Role:         aws.String(iamRole),
		Runtime:      aws.String(d.Get("runtime").(string)),
		Timeout:      aws.Int64(int64(d.Get("timeout").(int))),
	}

	if v, ok := d.GetOk("vpc_config"); ok {
		config, err := validateVPCConfig(v)
		if err != nil {
			return err
		}

		var subnetIds []*string
		for _, id := range config["subnet_ids"].(*schema.Set).List() {
			subnetIds = append(subnetIds, aws.String(id.(string)))
		}

		var securityGroupIds []*string
		for _, id := range config["security_group_ids"].(*schema.Set).List() {
			securityGroupIds = append(securityGroupIds, aws.String(id.(string)))
		}

		params.VpcConfig = &lambda.VpcConfig{
			SubnetIds:        subnetIds,
			SecurityGroupIds: securityGroupIds,
		}
	}

	// IAM profiles can take ~10 seconds to propagate in AWS:
	// http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#launch-instance-with-role-console
	// Error creating Lambda function: InvalidParameterValueException: The role defined for the task cannot be assumed by Lambda.
	err := resource.Retry(1*time.Minute, func() *resource.RetryError {
		_, err := conn.CreateFunction(params)
		if err != nil {
			log.Printf("[ERROR] Received %q, retrying CreateFunction", err)
			if awserr, ok := err.(awserr.Error); ok {
				if awserr.Code() == "InvalidParameterValueException" {
					log.Printf("[DEBUG] InvalidParameterValueException creating Lambda Function: %s", awserr)
					return resource.RetryableError(awserr)
				}
			}
			log.Printf("[DEBUG] Error creating Lambda Function: %s", err)
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("Error creating Lambda function: %s", err)
	}

	d.SetId(d.Get("function_name").(string))

	return resourceAwsLambdaFunctionRead(d, meta)
}

// resourceAwsLambdaFunctionRead maps to:
// GetFunction in the API / SDK
func resourceAwsLambdaFunctionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	log.Printf("[DEBUG] Fetching Lambda Function: %s", d.Id())

	params := &lambda.GetFunctionInput{
		FunctionName: aws.String(d.Get("function_name").(string)),
	}

	getFunctionOutput, err := conn.GetFunction(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "ResourceNotFoundException" {
			d.SetId("")
			return nil
		}
		return err
	}

	// getFunctionOutput.Code.Location is a pre-signed URL pointing at the zip
	// file that we uploaded when we created the resource. You can use it to
	// download the code from AWS. The other part is
	// getFunctionOutput.Configuration which holds metadata.

	function := getFunctionOutput.Configuration
	// TODO error checking / handling on the Set() calls.
	d.Set("arn", function.FunctionArn)
	d.Set("description", function.Description)
	d.Set("handler", function.Handler)
	d.Set("memory_size", function.MemorySize)
	d.Set("last_modified", function.LastModified)
	d.Set("role", function.Role)
	d.Set("runtime", function.Runtime)
	d.Set("timeout", function.Timeout)
	if config := flattenLambdaVpcConfigResponse(function.VpcConfig); len(config) > 0 {
		d.Set("vpc_config", config)
	}
	d.Set("source_code_hash", function.CodeSha256)

	return nil
}

// resourceAwsLambdaFunction maps to:
// DeleteFunction in the API / SDK
func resourceAwsLambdaFunctionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

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

	params := &lambda.DeleteFunctionInput{
		FunctionName: aws.String(d.Get("function_name").(string)),
	}

	_, err := conn.DeleteFunction(params)
	if err != nil {
		return fmt.Errorf("Error deleting Lambda Function: %s", err)
	}

	d.SetId("")

	return nil
}

// resourceAwsLambdaFunctionUpdate maps to:
// UpdateFunctionCode in the API / SDK
func resourceAwsLambdaFunctionUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	d.Partial(true)

	codeReq := &lambda.UpdateFunctionCodeInput{
		FunctionName: aws.String(d.Id()),
	}

	codeUpdate := false
	if v, ok := d.GetOk("filename"); ok && d.HasChange("source_code_hash") {
		file, err := loadFileContent(v.(string))
		if err != nil {
			return fmt.Errorf("Unable to load %q: %s", v.(string), err)
		}
		codeReq.ZipFile = file
		codeUpdate = true
	}
	if d.HasChange("s3_bucket") || d.HasChange("s3_key") || d.HasChange("s3_object_version") {
		codeReq.S3Bucket = aws.String(d.Get("s3_bucket").(string))
		codeReq.S3Key = aws.String(d.Get("s3_key").(string))
		codeReq.S3ObjectVersion = aws.String(d.Get("s3_object_version").(string))
		codeUpdate = true
	}

	log.Printf("[DEBUG] Send Update Lambda Function Code request: %#v", codeReq)
	if codeUpdate {
		_, err := conn.UpdateFunctionCode(codeReq)
		if err != nil {
			return fmt.Errorf("Error modifying Lambda Function Code %s: %s", d.Id(), err)
		}

		d.SetPartial("filename")
		d.SetPartial("source_code_hash")
		d.SetPartial("s3_bucket")
		d.SetPartial("s3_key")
		d.SetPartial("s3_object_version")
	}

	configReq := &lambda.UpdateFunctionConfigurationInput{
		FunctionName: aws.String(d.Id()),
	}

	configUpdate := false
	if d.HasChange("description") {
		configReq.Description = aws.String(d.Get("description").(string))
		configUpdate = true
	}
	if d.HasChange("handler") {
		configReq.Handler = aws.String(d.Get("handler").(string))
		configUpdate = true
	}
	if d.HasChange("memory_size") {
		configReq.MemorySize = aws.Int64(int64(d.Get("memory_size").(int)))
		configUpdate = true
	}
	if d.HasChange("role") {
		configReq.Role = aws.String(d.Get("role").(string))
		configUpdate = true
	}
	if d.HasChange("timeout") {
		configReq.Timeout = aws.Int64(int64(d.Get("timeout").(int)))
		configUpdate = true
	}

	log.Printf("[DEBUG] Send Update Lambda Function Configuration request: %#v", configReq)
	if configUpdate {
		_, err := conn.UpdateFunctionConfiguration(configReq)
		if err != nil {
			return fmt.Errorf("Error modifying Lambda Function Configuration %s: %s", d.Id(), err)
		}
		d.SetPartial("description")
		d.SetPartial("handler")
		d.SetPartial("memory_size")
		d.SetPartial("role")
		d.SetPartial("timeout")
	}
	d.Partial(false)

	return resourceAwsLambdaFunctionRead(d, meta)
}

// loadFileContent returns contents of a file in a given path
func loadFileContent(v string) ([]byte, error) {
	filename, err := homedir.Expand(v)
	if err != nil {
		return nil, err
	}
	fileContent, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return fileContent, nil
}

func validateVPCConfig(v interface{}) (map[string]interface{}, error) {
	configs := v.([]interface{})
	if len(configs) > 1 {
		return nil, errors.New("Only a single vpc_config block is expected")
	}

	config, ok := configs[0].(map[string]interface{})

	if !ok {
		return nil, errors.New("vpc_config is <nil>")
	}

	if config["subnet_ids"].(*schema.Set).Len() == 0 {
		return nil, errors.New("vpc_config.subnet_ids cannot be empty")
	}

	if config["security_group_ids"].(*schema.Set).Len() == 0 {
		return nil, errors.New("vpc_config.security_group_ids cannot be empty")
	}

	return config, nil
}
<commit_msg>provider/aws: Read VPC ID for Lambda function back from API<commit_after>package aws

import (
	"fmt"
	"io/ioutil"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/mitchellh/go-homedir"

	"errors"

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

func resourceAwsLambdaFunction() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLambdaFunctionCreate,
		Read:   resourceAwsLambdaFunctionRead,
		Update: resourceAwsLambdaFunctionUpdate,
		Delete: resourceAwsLambdaFunctionDelete,

		Schema: map[string]*schema.Schema{
			"filename": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"s3_bucket", "s3_key", "s3_object_version"},
			},
			"s3_bucket": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"s3_key": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"s3_object_version": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"function_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"handler": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"memory_size": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  128,
			},
			"role": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"runtime": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  "nodejs",
			},
			"timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  3,
			},
			"vpc_config": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"subnet_ids": &schema.Schema{
							Type:     schema.TypeSet,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set:      schema.HashString,
						},
						"security_group_ids": &schema.Schema{
							Type:     schema.TypeSet,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set:      schema.HashString,
						},
						"vpc_id": &schema.Schema{
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
			},
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"last_modified": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"source_code_hash": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
		},
	}
}

// resourceAwsLambdaFunction maps to:
// CreateFunction in the API / SDK
func resourceAwsLambdaFunctionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	functionName := d.Get("function_name").(string)
	iamRole := d.Get("role").(string)

	log.Printf("[DEBUG] Creating Lambda Function %s with role %s", functionName, iamRole)

	var functionCode *lambda.FunctionCode
	if v, ok := d.GetOk("filename"); ok {
		file, err := loadFileContent(v.(string))
		if err != nil {
			return fmt.Errorf("Unable to load %q: %s", v.(string), err)
		}
		functionCode = &lambda.FunctionCode{
			ZipFile: file,
		}
	} else {
		s3Bucket, bucketOk := d.GetOk("s3_bucket")
		s3Key, keyOk := d.GetOk("s3_key")
		s3ObjectVersion, versionOk := d.GetOk("s3_object_version")
		if !bucketOk || !keyOk {
			return errors.New("s3_bucket and s3_key must all be set while using S3 code source")
		}
		functionCode = &lambda.FunctionCode{
			S3Bucket: aws.String(s3Bucket.(string)),
			S3Key:    aws.String(s3Key.(string)),
		}
		if versionOk {
			functionCode.S3ObjectVersion = aws.String(s3ObjectVersion.(string))
		}
	}

	params := &lambda.CreateFunctionInput{
		Code:         functionCode,
		Description:  aws.String(d.Get("description").(string)),
		FunctionName: aws.String(functionName),
		Handler:      aws.String(d.Get("handler").(string)),
		MemorySize:   aws.Int64(int64(d.Get("memory_size").(int))),
		Role:         aws.String(iamRole),
		Runtime:      aws.String(d.Get("runtime").(string)),
		Timeout:      aws.Int64(int64(d.Get("timeout").(int))),
	}

	if v, ok := d.GetOk("vpc_config"); ok {
		config, err := validateVPCConfig(v)
		if err != nil {
			return err
		}

		var subnetIds []*string
		for _, id := range config["subnet_ids"].(*schema.Set).List() {
			subnetIds = append(subnetIds, aws.String(id.(string)))
		}

		var securityGroupIds []*string
		for _, id := range config["security_group_ids"].(*schema.Set).List() {
			securityGroupIds = append(securityGroupIds, aws.String(id.(string)))
		}

		params.VpcConfig = &lambda.VpcConfig{
			SubnetIds:        subnetIds,
			SecurityGroupIds: securityGroupIds,
		}
	}

	// IAM profiles can take ~10 seconds to propagate in AWS:
	// http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#launch-instance-with-role-console
	// Error creating Lambda function: InvalidParameterValueException: The role defined for the task cannot be assumed by Lambda.
	err := resource.Retry(1*time.Minute, func() *resource.RetryError {
		_, err := conn.CreateFunction(params)
		if err != nil {
			log.Printf("[ERROR] Received %q, retrying CreateFunction", err)
			if awserr, ok := err.(awserr.Error); ok {
				if awserr.Code() == "InvalidParameterValueException" {
					log.Printf("[DEBUG] InvalidParameterValueException creating Lambda Function: %s", awserr)
					return resource.RetryableError(awserr)
				}
			}
			log.Printf("[DEBUG] Error creating Lambda Function: %s", err)
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("Error creating Lambda function: %s", err)
	}

	d.SetId(d.Get("function_name").(string))

	return resourceAwsLambdaFunctionRead(d, meta)
}

// resourceAwsLambdaFunctionRead maps to:
// GetFunction in the API / SDK
func resourceAwsLambdaFunctionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	log.Printf("[DEBUG] Fetching Lambda Function: %s", d.Id())

	params := &lambda.GetFunctionInput{
		FunctionName: aws.String(d.Get("function_name").(string)),
	}

	getFunctionOutput, err := conn.GetFunction(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "ResourceNotFoundException" {
			d.SetId("")
			return nil
		}
		return err
	}

	// getFunctionOutput.Code.Location is a pre-signed URL pointing at the zip
	// file that we uploaded when we created the resource. You can use it to
	// download the code from AWS. The other part is
	// getFunctionOutput.Configuration which holds metadata.

	function := getFunctionOutput.Configuration
	// TODO error checking / handling on the Set() calls.
	d.Set("arn", function.FunctionArn)
	d.Set("description", function.Description)
	d.Set("handler", function.Handler)
	d.Set("memory_size", function.MemorySize)
	d.Set("last_modified", function.LastModified)
	d.Set("role", function.Role)
	d.Set("runtime", function.Runtime)
	d.Set("timeout", function.Timeout)
	if config := flattenLambdaVpcConfigResponse(function.VpcConfig); len(config) > 0 {
		log.Printf("[INFO] Setting Lambda %s VPC config %#v from API", d.Id(), config)
		err := d.Set("vpc_config", config)
		if err != nil {
			return fmt.Errorf("Failed setting vpc_config: %s", err)
		}
	}
	d.Set("source_code_hash", function.CodeSha256)

	return nil
}

// resourceAwsLambdaFunction maps to:
// DeleteFunction in the API / SDK
func resourceAwsLambdaFunctionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

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

	params := &lambda.DeleteFunctionInput{
		FunctionName: aws.String(d.Get("function_name").(string)),
	}

	_, err := conn.DeleteFunction(params)
	if err != nil {
		return fmt.Errorf("Error deleting Lambda Function: %s", err)
	}

	d.SetId("")

	return nil
}

// resourceAwsLambdaFunctionUpdate maps to:
// UpdateFunctionCode in the API / SDK
func resourceAwsLambdaFunctionUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	d.Partial(true)

	codeReq := &lambda.UpdateFunctionCodeInput{
		FunctionName: aws.String(d.Id()),
	}

	codeUpdate := false
	if v, ok := d.GetOk("filename"); ok && d.HasChange("source_code_hash") {
		file, err := loadFileContent(v.(string))
		if err != nil {
			return fmt.Errorf("Unable to load %q: %s", v.(string), err)
		}
		codeReq.ZipFile = file
		codeUpdate = true
	}
	if d.HasChange("s3_bucket") || d.HasChange("s3_key") || d.HasChange("s3_object_version") {
		codeReq.S3Bucket = aws.String(d.Get("s3_bucket").(string))
		codeReq.S3Key = aws.String(d.Get("s3_key").(string))
		codeReq.S3ObjectVersion = aws.String(d.Get("s3_object_version").(string))
		codeUpdate = true
	}

	log.Printf("[DEBUG] Send Update Lambda Function Code request: %#v", codeReq)
	if codeUpdate {
		_, err := conn.UpdateFunctionCode(codeReq)
		if err != nil {
			return fmt.Errorf("Error modifying Lambda Function Code %s: %s", d.Id(), err)
		}

		d.SetPartial("filename")
		d.SetPartial("source_code_hash")
		d.SetPartial("s3_bucket")
		d.SetPartial("s3_key")
		d.SetPartial("s3_object_version")
	}

	configReq := &lambda.UpdateFunctionConfigurationInput{
		FunctionName: aws.String(d.Id()),
	}

	configUpdate := false
	if d.HasChange("description") {
		configReq.Description = aws.String(d.Get("description").(string))
		configUpdate = true
	}
	if d.HasChange("handler") {
		configReq.Handler = aws.String(d.Get("handler").(string))
		configUpdate = true
	}
	if d.HasChange("memory_size") {
		configReq.MemorySize = aws.Int64(int64(d.Get("memory_size").(int)))
		configUpdate = true
	}
	if d.HasChange("role") {
		configReq.Role = aws.String(d.Get("role").(string))
		configUpdate = true
	}
	if d.HasChange("timeout") {
		configReq.Timeout = aws.Int64(int64(d.Get("timeout").(int)))
		configUpdate = true
	}

	log.Printf("[DEBUG] Send Update Lambda Function Configuration request: %#v", configReq)
	if configUpdate {
		_, err := conn.UpdateFunctionConfiguration(configReq)
		if err != nil {
			return fmt.Errorf("Error modifying Lambda Function Configuration %s: %s", d.Id(), err)
		}
		d.SetPartial("description")
		d.SetPartial("handler")
		d.SetPartial("memory_size")
		d.SetPartial("role")
		d.SetPartial("timeout")
	}
	d.Partial(false)

	return resourceAwsLambdaFunctionRead(d, meta)
}

// loadFileContent returns contents of a file in a given path
func loadFileContent(v string) ([]byte, error) {
	filename, err := homedir.Expand(v)
	if err != nil {
		return nil, err
	}
	fileContent, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return fileContent, nil
}

func validateVPCConfig(v interface{}) (map[string]interface{}, error) {
	configs := v.([]interface{})
	if len(configs) > 1 {
		return nil, errors.New("Only a single vpc_config block is expected")
	}

	config, ok := configs[0].(map[string]interface{})

	if !ok {
		return nil, errors.New("vpc_config is <nil>")
	}

	if config["subnet_ids"].(*schema.Set).Len() == 0 {
		return nil, errors.New("vpc_config.subnet_ids cannot be empty")
	}

	if config["security_group_ids"].(*schema.Set).Len() == 0 {
		return nil, errors.New("vpc_config.security_group_ids cannot be empty")
	}

	return config, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"github.com/codegangsta/cli"
	"io"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"text/template"
	"unicode"
)

const (
	TK_SYMBOL = iota
	TK_STRUCT_BEGIN
	TK_STRUCT_END
	TK_DATA_TYPE
	TK_ARRAY
	TK_EOF
)

var (
	datatypes map[string]map[string]struct {
		T string `json:"t"` // type
		R string `json:"r"` // read
		W string `json:"w"` // write
	} // type -> language -> t/r/w
)

var (
	TOKEN_EOF = &token{typ: TK_EOF}
)

type (
	field_info struct {
		Name  string
		Typ   string
		Array bool
	}
	struct_info struct {
		Name   string
		Fields []field_info
	}
)

type token struct {
	typ     int
	literal string
	r       rune
}

func syntax_error(p *Parser) {
	log.Println("syntax error @line:", p.lexer.lineno)
	log.Println(">> \033[1;31m", p.lexer.lines[p.lexer.lineno-1], "\033[0m <<")
	os.Exit(-1)
}

type Lexer struct {
	reader *bytes.Buffer
	lines  []string
	lineno int
}

func (lex *Lexer) init(r io.Reader) {
	bts, err := ioutil.ReadAll(r)
	if err != nil {
		log.Fatal(err)
	}

	// 按行读入源码
	scanner := bufio.NewScanner(bytes.NewBuffer(bts))
	for scanner.Scan() {
		lex.lines = append(lex.lines, scanner.Text())
	}

	// 清除注释
	re := regexp.MustCompile("(?m:^#(.*)$)")
	bts = re.ReplaceAllLiteral(bts, nil)
	lex.reader = bytes.NewBuffer(bts)
	lex.lineno = 1
}

func (lex *Lexer) next() (t *token) {
	defer func() {
		//log.Println(t)
	}()
	var r rune
	var err error
	for {
		r, _, err = lex.reader.ReadRune()
		if err == io.EOF {
			return TOKEN_EOF
		} else if unicode.IsSpace(r) {
			if r == '\n' {
				lex.lineno++
			}
			continue
		}
		break
	}

	if r == '=' {
		for k := 0; k < 2; k++ { // check "==="
			r, _, err = lex.reader.ReadRune()
			if err == io.EOF {
				return TOKEN_EOF
			}
			if r != '=' {
				lex.reader.UnreadRune()
				return &token{typ: TK_STRUCT_BEGIN}
			}
		}
		return &token{typ: TK_STRUCT_END}
	} else if unicode.IsLetter(r) {
		var runes []rune
		for {
			runes = append(runes, r)
			r, _, err = lex.reader.ReadRune()
			if err == io.EOF {
				break
			} else if unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' {
				continue
			} else {
				lex.reader.UnreadRune()
				break
			}
		}

		t := &token{}
		t.literal = string(runes)
		if _, ok := datatypes[t.literal]; ok {
			t.typ = TK_DATA_TYPE
		} else if t.literal == "array" {
			t.typ = TK_ARRAY
		} else {
			t.typ = TK_SYMBOL
		}

		return t
	} else {
		log.Fatal("lex error @line:", lex.lineno)
	}
	return nil
}

func (lex *Lexer) eof() bool {
	for {
		r, _, err := lex.reader.ReadRune()
		if err == io.EOF {
			return true
		} else if unicode.IsSpace(r) {
			if r == '\n' {
				lex.lineno++
			}
			continue
		} else {
			lex.reader.UnreadRune()
			return false
		}
	}
}

//////////////////////////////////////////////////////////////
type Parser struct {
	lexer   *Lexer
	infos   []struct_info
	symbols map[string]bool
}

func (p *Parser) init(lex *Lexer) {
	p.lexer = lex
	p.symbols = make(map[string]bool)
}

func (p *Parser) match(typ int) *token {
	t := p.lexer.next()
	if t.typ != typ {
		syntax_error(p)
	}
	return t
}

func (p *Parser) expr() bool {
	if p.lexer.eof() {
		return false
	}
	info := struct_info{}

	t := p.match(TK_SYMBOL)
	info.Name = t.literal
	p.symbols[t.literal] = true
	p.match(TK_STRUCT_BEGIN)
	p.fields(&info)
	p.infos = append(p.infos, info)
	return true
}

func (p *Parser) fields(info *struct_info) {
	for {
		t := p.lexer.next()
		if t.typ == TK_STRUCT_END {
			return
		}
		if t.typ != TK_SYMBOL {
			syntax_error(p)
		}

		field := field_info{Name: t.literal}
		t = p.lexer.next()
		if t.typ == TK_ARRAY {
			field.Array = true
			t = p.lexer.next()
		}

		if t.typ == TK_DATA_TYPE || t.typ == TK_SYMBOL {
			field.Typ = t.literal
		} else {
			syntax_error(p)
		}

		info.Fields = append(info.Fields, field)
	}
}

func (p *Parser) semantic_check() {
	for _, info := range p.infos {
	FIELDLOOP:
		for _, field := range info.Fields {
			if _, ok := datatypes[field.Typ]; !ok {
				if p.symbols[field.Typ] {
					continue FIELDLOOP
				}
				log.Fatal("symbol not found:", field)
			}
		}
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "Protocol Data Structure Generator"
	app.Usage = "handle proto.txt"
	app.Authors = []cli.Author{{Name: "xtaci"}, {Name: "ycs"}}
	app.Version = "1.0"
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "file,f", Value: "./proto.txt", Usage: "input proto.txt file"},
		cli.StringFlag{Name: "binding,b", Value: "go", Usage: "language type binding"},
		cli.StringFlag{Name: "template,t", Value: "./templates/server/proto.tmpl", Usage: "template file"},
	}
	app.Action = func(c *cli.Context) {
		// load primitives mapping
		f, err := os.Open("primitives.json")
		if err != nil {
			log.Fatal(err)
		}
		if err := json.NewDecoder(f).Decode(&datatypes); err != nil {
			log.Fatal(err)
		}

		// parse
		file, err := os.Open(c.String("file"))
		if err != nil {
			log.Fatal(err)
		}
		lexer := Lexer{}
		lexer.init(file)
		p := Parser{}
		p.init(&lexer)
		for p.expr() {
		}

		// semantic
		p.semantic_check()

		// use template to generate final output
		funcMap := template.FuncMap{
			"Type": func(t string) string {
				return datatypes[t][c.String("binding")].T
			},
			"Read": func(t string) string {
				return datatypes[t][c.String("binding")].R
			},
			"Write": func(t string) string {
				return datatypes[t][c.String("binding")].W
			},
		}
		tmpl, err := template.New("proto.tmpl").Funcs(funcMap).ParseFiles(c.String("template"))
		if err != nil {
			log.Fatal(err)
		}
		err = tmpl.Execute(os.Stdout, p.infos)
		if err != nil {
			log.Fatal(err)
		}
	}
	app.Run(os.Args)
}
<commit_msg>update usage<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"github.com/codegangsta/cli"
	"io"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"text/template"
	"unicode"
)

const (
	TK_SYMBOL = iota
	TK_STRUCT_BEGIN
	TK_STRUCT_END
	TK_DATA_TYPE
	TK_ARRAY
	TK_EOF
)

var (
	datatypes map[string]map[string]struct {
		T string `json:"t"` // type
		R string `json:"r"` // read
		W string `json:"w"` // write
	} // type -> language -> t/r/w
)

var (
	TOKEN_EOF = &token{typ: TK_EOF}
)

type (
	field_info struct {
		Name  string
		Typ   string
		Array bool
	}
	struct_info struct {
		Name   string
		Fields []field_info
	}
)

type token struct {
	typ     int
	literal string
	r       rune
}

func syntax_error(p *Parser) {
	log.Println("syntax error @line:", p.lexer.lineno)
	log.Println(">> \033[1;31m", p.lexer.lines[p.lexer.lineno-1], "\033[0m <<")
	os.Exit(-1)
}

type Lexer struct {
	reader *bytes.Buffer
	lines  []string
	lineno int
}

func (lex *Lexer) init(r io.Reader) {
	bts, err := ioutil.ReadAll(r)
	if err != nil {
		log.Fatal(err)
	}

	// 按行读入源码
	scanner := bufio.NewScanner(bytes.NewBuffer(bts))
	for scanner.Scan() {
		lex.lines = append(lex.lines, scanner.Text())
	}

	// 清除注释
	re := regexp.MustCompile("(?m:^#(.*)$)")
	bts = re.ReplaceAllLiteral(bts, nil)
	lex.reader = bytes.NewBuffer(bts)
	lex.lineno = 1
}

func (lex *Lexer) next() (t *token) {
	defer func() {
		//log.Println(t)
	}()
	var r rune
	var err error
	for {
		r, _, err = lex.reader.ReadRune()
		if err == io.EOF {
			return TOKEN_EOF
		} else if unicode.IsSpace(r) {
			if r == '\n' {
				lex.lineno++
			}
			continue
		}
		break
	}

	if r == '=' {
		for k := 0; k < 2; k++ { // check "==="
			r, _, err = lex.reader.ReadRune()
			if err == io.EOF {
				return TOKEN_EOF
			}
			if r != '=' {
				lex.reader.UnreadRune()
				return &token{typ: TK_STRUCT_BEGIN}
			}
		}
		return &token{typ: TK_STRUCT_END}
	} else if unicode.IsLetter(r) {
		var runes []rune
		for {
			runes = append(runes, r)
			r, _, err = lex.reader.ReadRune()
			if err == io.EOF {
				break
			} else if unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' {
				continue
			} else {
				lex.reader.UnreadRune()
				break
			}
		}

		t := &token{}
		t.literal = string(runes)
		if _, ok := datatypes[t.literal]; ok {
			t.typ = TK_DATA_TYPE
		} else if t.literal == "array" {
			t.typ = TK_ARRAY
		} else {
			t.typ = TK_SYMBOL
		}

		return t
	} else {
		log.Fatal("lex error @line:", lex.lineno)
	}
	return nil
}

func (lex *Lexer) eof() bool {
	for {
		r, _, err := lex.reader.ReadRune()
		if err == io.EOF {
			return true
		} else if unicode.IsSpace(r) {
			if r == '\n' {
				lex.lineno++
			}
			continue
		} else {
			lex.reader.UnreadRune()
			return false
		}
	}
}

//////////////////////////////////////////////////////////////
type Parser struct {
	lexer   *Lexer
	infos   []struct_info
	symbols map[string]bool
}

func (p *Parser) init(lex *Lexer) {
	p.lexer = lex
	p.symbols = make(map[string]bool)
}

func (p *Parser) match(typ int) *token {
	t := p.lexer.next()
	if t.typ != typ {
		syntax_error(p)
	}
	return t
}

func (p *Parser) expr() bool {
	if p.lexer.eof() {
		return false
	}
	info := struct_info{}

	t := p.match(TK_SYMBOL)
	info.Name = t.literal
	p.symbols[t.literal] = true
	p.match(TK_STRUCT_BEGIN)
	p.fields(&info)
	p.infos = append(p.infos, info)
	return true
}

func (p *Parser) fields(info *struct_info) {
	for {
		t := p.lexer.next()
		if t.typ == TK_STRUCT_END {
			return
		}
		if t.typ != TK_SYMBOL {
			syntax_error(p)
		}

		field := field_info{Name: t.literal}
		t = p.lexer.next()
		if t.typ == TK_ARRAY {
			field.Array = true
			t = p.lexer.next()
		}

		if t.typ == TK_DATA_TYPE || t.typ == TK_SYMBOL {
			field.Typ = t.literal
		} else {
			syntax_error(p)
		}

		info.Fields = append(info.Fields, field)
	}
}

func (p *Parser) semantic_check() {
	for _, info := range p.infos {
	FIELDLOOP:
		for _, field := range info.Fields {
			if _, ok := datatypes[field.Typ]; !ok {
				if p.symbols[field.Typ] {
					continue FIELDLOOP
				}
				log.Fatal("symbol not found:", field)
			}
		}
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "Protocol Data Structure Generator"
	app.Usage = "handle proto.txt"
	app.Authors = []cli.Author{{Name: "xtaci"}, {Name: "ycs"}}
	app.Version = "1.0"
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "file,f", Value: "./proto.txt", Usage: "input proto.txt file"},
		cli.StringFlag{Name: "binding,b", Value: "go", Usage: `language type binding:"go","cs"`},
		cli.StringFlag{Name: "template,t", Value: "./templates/server/proto.tmpl", Usage: "template file"},
	}
	app.Action = func(c *cli.Context) {
		// load primitives mapping
		f, err := os.Open("primitives.json")
		if err != nil {
			log.Fatal(err)
		}
		if err := json.NewDecoder(f).Decode(&datatypes); err != nil {
			log.Fatal(err)
		}

		// parse
		file, err := os.Open(c.String("file"))
		if err != nil {
			log.Fatal(err)
		}
		lexer := Lexer{}
		lexer.init(file)
		p := Parser{}
		p.init(&lexer)
		for p.expr() {
		}

		// semantic
		p.semantic_check()

		// use template to generate final output
		funcMap := template.FuncMap{
			"Type": func(t string) string {
				return datatypes[t][c.String("binding")].T
			},
			"Read": func(t string) string {
				return datatypes[t][c.String("binding")].R
			},
			"Write": func(t string) string {
				return datatypes[t][c.String("binding")].W
			},
		}
		tmpl, err := template.New("proto.tmpl").Funcs(funcMap).ParseFiles(c.String("template"))
		if err != nil {
			log.Fatal(err)
		}
		err = tmpl.Execute(os.Stdout, p.infos)
		if err != nil {
			log.Fatal(err)
		}
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin dragonfly freebsd linux netbsd openbsd solaris

package signal

import (
	"flag"
	"io/ioutil"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"syscall"
	"testing"
	"time"
)

func waitSig(t *testing.T, c <-chan os.Signal, sig os.Signal) {
	select {
	case s := <-c:
		if s != sig {
			t.Fatalf("signal was %v, want %v", s, sig)
		}
	case <-time.After(1 * time.Second):
		t.Fatalf("timeout waiting for %v", sig)
	}
}

// Test that basic signal handling works.
func TestSignal(t *testing.T) {
	// Ask for SIGHUP
	c := make(chan os.Signal, 1)
	Notify(c, syscall.SIGHUP)
	defer Stop(c)

	// Send this process a SIGHUP
	t.Logf("sighup...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c, syscall.SIGHUP)

	// Ask for everything we can get.
	c1 := make(chan os.Signal, 1)
	Notify(c1)

	// Send this process a SIGWINCH
	t.Logf("sigwinch...")
	syscall.Kill(syscall.Getpid(), syscall.SIGWINCH)
	waitSig(t, c1, syscall.SIGWINCH)

	// Send two more SIGHUPs, to make sure that
	// they get delivered on c1 and that not reading
	// from c does not block everything.
	t.Logf("sighup...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c1, syscall.SIGHUP)
	t.Logf("sighup...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c1, syscall.SIGHUP)

	// The first SIGHUP should be waiting for us on c.
	waitSig(t, c, syscall.SIGHUP)
}

func TestStress(t *testing.T) {
	dur := 3 * time.Second
	if testing.Short() {
		dur = 100 * time.Millisecond
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(4))
	done := make(chan bool)
	finished := make(chan bool)
	go func() {
		sig := make(chan os.Signal, 1)
		Notify(sig, syscall.SIGUSR1)
		defer Stop(sig)
	Loop:
		for {
			select {
			case <-sig:
			case <-done:
				break Loop
			}
		}
		finished <- true
	}()
	go func() {
	Loop:
		for {
			select {
			case <-done:
				break Loop
			default:
				syscall.Kill(syscall.Getpid(), syscall.SIGUSR1)
				runtime.Gosched()
			}
		}
		finished <- true
	}()
	time.Sleep(dur)
	close(done)
	<-finished
	<-finished
	// When run with 'go test -cpu=1,2,4' SIGUSR1 from this test can slip
	// into subsequent TestSignal() causing failure.
	// Sleep for a while to reduce the possibility of the failure.
	time.Sleep(10 * time.Millisecond)
}

var sendUncaughtSighup = flag.Int("send_uncaught_sighup", 0, "send uncaught SIGHUP during TestStop")

// Test that Stop cancels the channel's registrations.
func TestStop(t *testing.T) {
	sigs := []syscall.Signal{
		syscall.SIGWINCH,
		syscall.SIGHUP,
	}

	for _, sig := range sigs {
		// Send the signal.
		// If it's SIGWINCH, we should not see it.
		// If it's SIGHUP, maybe we'll die. Let the flag tell us what to do.
		if sig != syscall.SIGHUP || *sendUncaughtSighup == 1 {
			syscall.Kill(syscall.Getpid(), sig)
		}
		time.Sleep(10 * time.Millisecond)

		// Ask for signal
		c := make(chan os.Signal, 1)
		Notify(c, sig)
		defer Stop(c)

		// Send this process that signal
		syscall.Kill(syscall.Getpid(), sig)
		waitSig(t, c, sig)

		Stop(c)
		select {
		case s := <-c:
			t.Fatalf("unexpected signal %v", s)
		case <-time.After(10 * time.Millisecond):
			// nothing to read - good
		}

		// Send the signal.
		// If it's SIGWINCH, we should not see it.
		// If it's SIGHUP, maybe we'll die. Let the flag tell us what to do.
		if sig != syscall.SIGHUP || *sendUncaughtSighup == 2 {
			syscall.Kill(syscall.Getpid(), sig)
		}

		select {
		case s := <-c:
			t.Fatalf("unexpected signal %v", s)
		case <-time.After(10 * time.Millisecond):
			// nothing to read - good
		}
	}
}

// Test that when run under nohup, an uncaught SIGHUP does not kill the program,
// but a
func TestNohup(t *testing.T) {
	// Ugly: ask for SIGHUP so that child will not have no-hup set
	// even if test is running under nohup environment.
	// We have no intention of reading from c.
	c := make(chan os.Signal, 1)
	Notify(c, syscall.SIGHUP)

	// When run without nohup, the test should crash on an uncaught SIGHUP.
	// When run under nohup, the test should ignore uncaught SIGHUPs,
	// because the runtime is not supposed to be listening for them.
	// Either way, TestStop should still be able to catch them when it wants them
	// and then when it stops wanting them, the original behavior should resume.
	//
	// send_uncaught_sighup=1 sends the SIGHUP before starting to listen for SIGHUPs.
	// send_uncaught_sighup=2 sends the SIGHUP after no longer listening for SIGHUPs.
	//
	// Both should fail without nohup and succeed with nohup.

	for i := 1; i <= 2; i++ {
		out, err := exec.Command(os.Args[0], "-test.run=TestStop", "-send_uncaught_sighup="+strconv.Itoa(i)).CombinedOutput()
		if err == nil {
			t.Fatalf("ran test with -send_uncaught_sighup=%d and it succeeded: expected failure.\nOutput:\n%s", i, out)
		}
	}

	Stop(c)

	// Again, this time with nohup, assuming we can find it.
	_, err := os.Stat("/usr/bin/nohup")
	if err != nil {
		t.Skip("cannot find nohup; skipping second half of test")
	}

	for i := 1; i <= 2; i++ {
		os.Remove("nohup.out")
		out, err := exec.Command("/usr/bin/nohup", os.Args[0], "-test.run=TestStop", "-send_uncaught_sighup="+strconv.Itoa(i)).CombinedOutput()

		data, _ := ioutil.ReadFile("nohup.out")
		os.Remove("nohup.out")
		if err != nil {
			t.Fatalf("ran test with -send_uncaught_sighup=%d under nohup and it failed: expected success.\nError: %v\nOutput:\n%s%s", i, err, out, data)
		}
	}
}
<commit_msg>os/signal: increase timeout from 10ms to 100ms<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin dragonfly freebsd linux netbsd openbsd solaris

package signal

import (
	"flag"
	"io/ioutil"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"syscall"
	"testing"
	"time"
)

func waitSig(t *testing.T, c <-chan os.Signal, sig os.Signal) {
	select {
	case s := <-c:
		if s != sig {
			t.Fatalf("signal was %v, want %v", s, sig)
		}
	case <-time.After(1 * time.Second):
		t.Fatalf("timeout waiting for %v", sig)
	}
}

// Test that basic signal handling works.
func TestSignal(t *testing.T) {
	// Ask for SIGHUP
	c := make(chan os.Signal, 1)
	Notify(c, syscall.SIGHUP)
	defer Stop(c)

	// Send this process a SIGHUP
	t.Logf("sighup...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c, syscall.SIGHUP)

	// Ask for everything we can get.
	c1 := make(chan os.Signal, 1)
	Notify(c1)

	// Send this process a SIGWINCH
	t.Logf("sigwinch...")
	syscall.Kill(syscall.Getpid(), syscall.SIGWINCH)
	waitSig(t, c1, syscall.SIGWINCH)

	// Send two more SIGHUPs, to make sure that
	// they get delivered on c1 and that not reading
	// from c does not block everything.
	t.Logf("sighup...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c1, syscall.SIGHUP)
	t.Logf("sighup...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c1, syscall.SIGHUP)

	// The first SIGHUP should be waiting for us on c.
	waitSig(t, c, syscall.SIGHUP)
}

func TestStress(t *testing.T) {
	dur := 3 * time.Second
	if testing.Short() {
		dur = 100 * time.Millisecond
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(4))
	done := make(chan bool)
	finished := make(chan bool)
	go func() {
		sig := make(chan os.Signal, 1)
		Notify(sig, syscall.SIGUSR1)
		defer Stop(sig)
	Loop:
		for {
			select {
			case <-sig:
			case <-done:
				break Loop
			}
		}
		finished <- true
	}()
	go func() {
	Loop:
		for {
			select {
			case <-done:
				break Loop
			default:
				syscall.Kill(syscall.Getpid(), syscall.SIGUSR1)
				runtime.Gosched()
			}
		}
		finished <- true
	}()
	time.Sleep(dur)
	close(done)
	<-finished
	<-finished
	// When run with 'go test -cpu=1,2,4' SIGUSR1 from this test can slip
	// into subsequent TestSignal() causing failure.
	// Sleep for a while to reduce the possibility of the failure.
	time.Sleep(10 * time.Millisecond)
}

var sendUncaughtSighup = flag.Int("send_uncaught_sighup", 0, "send uncaught SIGHUP during TestStop")

// Test that Stop cancels the channel's registrations.
func TestStop(t *testing.T) {
	sigs := []syscall.Signal{
		syscall.SIGWINCH,
		syscall.SIGHUP,
	}

	for _, sig := range sigs {
		// Send the signal.
		// If it's SIGWINCH, we should not see it.
		// If it's SIGHUP, maybe we'll die. Let the flag tell us what to do.
		if sig != syscall.SIGHUP || *sendUncaughtSighup == 1 {
			syscall.Kill(syscall.Getpid(), sig)
		}
		time.Sleep(100 * time.Millisecond)

		// Ask for signal
		c := make(chan os.Signal, 1)
		Notify(c, sig)
		defer Stop(c)

		// Send this process that signal
		syscall.Kill(syscall.Getpid(), sig)
		waitSig(t, c, sig)

		Stop(c)
		select {
		case s := <-c:
			t.Fatalf("unexpected signal %v", s)
		case <-time.After(100 * time.Millisecond):
			// nothing to read - good
		}

		// Send the signal.
		// If it's SIGWINCH, we should not see it.
		// If it's SIGHUP, maybe we'll die. Let the flag tell us what to do.
		if sig != syscall.SIGHUP || *sendUncaughtSighup == 2 {
			syscall.Kill(syscall.Getpid(), sig)
		}

		select {
		case s := <-c:
			t.Fatalf("unexpected signal %v", s)
		case <-time.After(100 * time.Millisecond):
			// nothing to read - good
		}
	}
}

// Test that when run under nohup, an uncaught SIGHUP does not kill the program,
// but a
func TestNohup(t *testing.T) {
	// Ugly: ask for SIGHUP so that child will not have no-hup set
	// even if test is running under nohup environment.
	// We have no intention of reading from c.
	c := make(chan os.Signal, 1)
	Notify(c, syscall.SIGHUP)

	// When run without nohup, the test should crash on an uncaught SIGHUP.
	// When run under nohup, the test should ignore uncaught SIGHUPs,
	// because the runtime is not supposed to be listening for them.
	// Either way, TestStop should still be able to catch them when it wants them
	// and then when it stops wanting them, the original behavior should resume.
	//
	// send_uncaught_sighup=1 sends the SIGHUP before starting to listen for SIGHUPs.
	// send_uncaught_sighup=2 sends the SIGHUP after no longer listening for SIGHUPs.
	//
	// Both should fail without nohup and succeed with nohup.

	for i := 1; i <= 2; i++ {
		out, err := exec.Command(os.Args[0], "-test.run=TestStop", "-send_uncaught_sighup="+strconv.Itoa(i)).CombinedOutput()
		if err == nil {
			t.Fatalf("ran test with -send_uncaught_sighup=%d and it succeeded: expected failure.\nOutput:\n%s", i, out)
		}
	}

	Stop(c)

	// Again, this time with nohup, assuming we can find it.
	_, err := os.Stat("/usr/bin/nohup")
	if err != nil {
		t.Skip("cannot find nohup; skipping second half of test")
	}

	for i := 1; i <= 2; i++ {
		os.Remove("nohup.out")
		out, err := exec.Command("/usr/bin/nohup", os.Args[0], "-test.run=TestStop", "-send_uncaught_sighup="+strconv.Itoa(i)).CombinedOutput()

		data, _ := ioutil.ReadFile("nohup.out")
		os.Remove("nohup.out")
		if err != nil {
			t.Fatalf("ran test with -send_uncaught_sighup=%d under nohup and it failed: expected success.\nError: %v\nOutput:\n%s%s", i, err, out, data)
		}
	}
}
<|endoftext|>


<commit_before>package tests

import (
	"fmt"
	"testing"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/portworx/torpedo/drivers/node"
	. "github.com/portworx/torpedo/tests"
)

func TestReboot(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Torpedo : Reboot")
}

var _ = BeforeSuite(func() {
	InitInstance()
})

func rebootNodesTest(testName string, allNodes bool) {
	It("has to scheduler apps and reboot app node(s)", func() {
		var err error
		contexts := ScheduleAndValidate(testName)

		Step("get nodes for all apps in test and reboot their nodes", func() {
			for _, ctx := range contexts {
				var appNodes []node.Node
				var nodesToReboot []node.Node

				Step(fmt.Sprintf("get nodes for %s app", ctx.App.Key), func() {
					appNodes, err = Inst().S.GetNodesForApp(ctx)
					Expect(err).NotTo(HaveOccurred())
					Expect(appNodes).NotTo(BeEmpty())
				})

				Step(fmt.Sprintf("find node(s) to reboot for %s app", ctx.App.Key), func() {
					if allNodes {
						nodesToReboot = appNodes
					} else {
						nodesToReboot = append(nodesToReboot, appNodes[0])
					}
				})

				Step(fmt.Sprintf("reboot app %s's node(s): %v", ctx.App.Key, nodesToReboot), func() {
					for _, n := range nodesToReboot {
						err = Inst().N.RebootNode(n, node.RebootNodeOpts{
							Force: false,
							ConnectionOpts: node.ConnectionOpts{
								Timeout:         1 * time.Minute,
								TimeBeforeRetry: 5 * time.Second,
							},
						})
						Expect(err).NotTo(HaveOccurred())

						Step("wait for node to go down", func() {
							time.Sleep(20 * time.Second)
						})

						Step("wait for node to be back up", func() {
							err = Inst().N.TestConnection(n, node.ConnectionOpts{
								Timeout:         15 * time.Minute,
								TimeBeforeRetry: 10 * time.Second,
							})
							Expect(err).NotTo(HaveOccurred())
						})

						Step(fmt.Sprintf("wait to scheduler: %s and volume driver: %s to start",
							Inst().S.String(), Inst().V.String()), func() {

							err = Inst().S.IsNodeReady(n)
							Expect(err).NotTo(HaveOccurred())

							err = Inst().V.WaitDriverUpOnNode(n)
							Expect(err).NotTo(HaveOccurred())
						})
					}
				})
			}
		})

		Step("validate and destroy apps", func() {
			for _, ctx := range contexts {
				ValidateAndDestroy(ctx, nil)
			}
		})

	})
}

var _ = Describe("Reboot one node test", func() {
	rebootNodesTest("rebootonenode", false)
})

var _ = Describe("Reboot all nodes test", func() {
	rebootNodesTest("rebootallnodes", false)
})

var _ = AfterSuite(func() {
	CollectSupport()
	ValidateCleanup()
})

func init() {
	ParseFlags()
}
<commit_msg>Reboot all node test should reboot all nodes<commit_after>package tests

import (
	"fmt"
	"testing"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/portworx/torpedo/drivers/node"
	. "github.com/portworx/torpedo/tests"
)

func TestReboot(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Torpedo : Reboot")
}

var _ = BeforeSuite(func() {
	InitInstance()
})

func rebootNodesTest(testName string, allNodes bool) {
	It("has to scheduler apps and reboot app node(s)", func() {
		var err error
		contexts := ScheduleAndValidate(testName)

		Step("get nodes for all apps in test and reboot their nodes", func() {
			for _, ctx := range contexts {
				var appNodes []node.Node
				var nodesToReboot []node.Node

				Step(fmt.Sprintf("get nodes for %s app", ctx.App.Key), func() {
					appNodes, err = Inst().S.GetNodesForApp(ctx)
					Expect(err).NotTo(HaveOccurred())
					Expect(appNodes).NotTo(BeEmpty())
				})

				Step(fmt.Sprintf("find node(s) to reboot for %s app", ctx.App.Key), func() {
					if allNodes {
						nodesToReboot = appNodes
					} else {
						nodesToReboot = append(nodesToReboot, appNodes[0])
					}
				})

				Step(fmt.Sprintf("reboot app %s's node(s): %v", ctx.App.Key, nodesToReboot), func() {
					for _, n := range nodesToReboot {
						err = Inst().N.RebootNode(n, node.RebootNodeOpts{
							Force: false,
							ConnectionOpts: node.ConnectionOpts{
								Timeout:         1 * time.Minute,
								TimeBeforeRetry: 5 * time.Second,
							},
						})
						Expect(err).NotTo(HaveOccurred())

						Step("wait for node to go down", func() {
							time.Sleep(20 * time.Second)
						})

						Step("wait for node to be back up", func() {
							err = Inst().N.TestConnection(n, node.ConnectionOpts{
								Timeout:         15 * time.Minute,
								TimeBeforeRetry: 10 * time.Second,
							})
							Expect(err).NotTo(HaveOccurred())
						})

						Step(fmt.Sprintf("wait to scheduler: %s and volume driver: %s to start",
							Inst().S.String(), Inst().V.String()), func() {

							err = Inst().S.IsNodeReady(n)
							Expect(err).NotTo(HaveOccurred())

							err = Inst().V.WaitDriverUpOnNode(n)
							Expect(err).NotTo(HaveOccurred())
						})
					}
				})
			}
		})

		Step("validate and destroy apps", func() {
			for _, ctx := range contexts {
				ValidateAndDestroy(ctx, nil)
			}
		})

	})
}

var _ = Describe("Reboot one node test", func() {
	rebootNodesTest("rebootonenode", false)
})

var _ = Describe("Reboot all nodes test", func() {
	rebootNodesTest("rebootallnodes", true)
})

var _ = AfterSuite(func() {
	CollectSupport()
	ValidateCleanup()
})

func init() {
	ParseFlags()
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package user

import (
	"fmt"
	"syscall"
	"unsafe"
)

func lookupFullName(domain, username, domainAndUser string) (string, error) {
	// try domain controller first
	name, e := syscall.TranslateAccountName(domainAndUser,
		syscall.NameSamCompatible, syscall.NameDisplay, 50)
	if e != nil {
		// domain lookup failed, perhaps this pc is not part of domain
		d := syscall.StringToUTF16Ptr(domain)
		u := syscall.StringToUTF16Ptr(username)
		var p *byte
		e := syscall.NetUserGetInfo(d, u, 10, &p)
		if e != nil {
			return "", e
		}
		defer syscall.NetApiBufferFree(p)
		i := (*syscall.UserInfo10)(unsafe.Pointer(p))
		if i.FullName == nil {
			return "", nil
		}
		name = syscall.UTF16ToString((*[1024]uint16)(unsafe.Pointer(i.FullName))[:])
	}
	return name, nil
}

func newUser(usid *syscall.SID, gid, dir string) (*User, error) {
	username, domain, t, e := usid.LookupAccount("")
	if e != nil {
		return nil, e
	}
	if t != syscall.SidTypeUser {
		return nil, fmt.Errorf("user: should be user account type, not %d", t)
	}
	domainAndUser := domain + `\` + username
	uid, e := usid.String()
	if e != nil {
		return nil, e
	}
	name, e := lookupFullName(domain, username, domainAndUser)
	if e != nil {
		return nil, e
	}
	u := &User{
		Uid:      uid,
		Gid:      gid,
		Username: domainAndUser,
		Name:     name,
		HomeDir:  dir,
	}
	return u, nil
}

// Current returns the current user.
func Current() (*User, error) {
	t, e := syscall.OpenCurrentProcessToken()
	if e != nil {
		return nil, e
	}
	u, e := t.GetTokenUser()
	if e != nil {
		return nil, e
	}
	pg, e := t.GetTokenPrimaryGroup()
	if e != nil {
		return nil, e
	}
	gid, e := pg.PrimaryGroup.String()
	if e != nil {
		return nil, e
	}
	dir, e := t.GetUserProfileDirectory()
	if e != nil {
		return nil, e
	}
	return newUser(u.User.Sid, gid, dir)
}

// BUG(brainman): Lookup and LookupId functions do not set
// Gid and HomeDir fields in the User struct returned on windows.

func newUserFromSid(usid *syscall.SID) (*User, error) {
	// TODO(brainman): do not know where to get gid and dir fields
	gid := "unknown"
	dir := "Unknown directory"
	return newUser(usid, gid, dir)
}

// Lookup looks up a user by username.
func Lookup(username string) (*User, error) {
	sid, _, t, e := syscall.LookupSID("", username)
	if e != nil {
		return nil, e
	}
	if t != syscall.SidTypeUser {
		return nil, fmt.Errorf("user: should be user account type, not %d", t)
	}
	return newUserFromSid(sid)
}

// LookupId looks up a user by userid.
func LookupId(uid string) (*User, error) {
	sid, e := syscall.StringToSid(uid)
	if e != nil {
		return nil, e
	}
	return newUserFromSid(sid)
}
<commit_msg>Deleted windoes file.<commit_after><|endoftext|>


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

// +build go1.3

package lxd

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/schema"
	"gopkg.in/juju/environschema.v1"

	"github.com/juju/juju/container/lxd/lxdclient"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/instance"
)

// TODO(ericsnow) Support providing cert/key file.

// The LXD-specific config keys.
const (
	cfgNamespace  = "namespace"
	cfgRemoteURL  = "remote-url"
	cfgClientCert = "client-cert"
	cfgClientKey  = "client-key"
)

// TODO(ericsnow) Use configSchema.ExampleYAML (once it is implemented)
// to generate boilerplaceConfig.

// boilerplateConfig will be shown in help output, so please keep it up to
// date when you change environment configuration below.
var boilerplateConfig = `
lxd:
    type: lxd

    # namespace identifies the namespace to associate with containers
    # created by the provider.  It is prepended to the container names.
    # By default the environment's name is used as the namespace.
    #
    # namespace: lxd

    # remote-url is the URL to the LXD API server to use for managing
    # containers, if any. If not specified then the locally running LXD
    # server is used.
    #
    # remote-url:

    # The cert and key the client should use to connect to the remote
    # may also be provided. If not then they are auto-generated.
    #
    # client-cert:
    # client-key:

`[1:]

// configSchema defines the schema for the configuration attributes
// defined by the LXD provider.
var configSchema = environschema.Fields{
	cfgNamespace: {
		Description: `Identifies the namespace to associate with containers created by the provider.  It is prepended to the container names.  By default the environment's name is used as the namespace.`,
		Type:        environschema.Tstring,
		Immutable:   true,
	},
	cfgRemoteURL: {
		Description: `Identifies the LXD API server to use for managing containers, if any.`,
		Type:        environschema.Tstring,
		Immutable:   true,
	},
	cfgClientKey: {
		Description: `The client key used for connecting to a LXD host machine.`,
		Type:        environschema.Tstring,
		Immutable:   true,
	},
	cfgClientCert: {
		Description: `The client cert used for connecting to a LXD host machine.`,
		Type:        environschema.Tstring,
		Immutable:   true,
	},
}

var (
	// TODO(ericsnow) Extract the defaults from configSchema as soon as
	// (or if) environschema.Attr supports defaults.

	configBaseDefaults = schema.Defaults{
		cfgNamespace:  "",
		cfgRemoteURL:  "",
		cfgClientCert: "",
		cfgClientKey:  "",
	}

	configFields, configDefaults = func() (schema.Fields, schema.Defaults) {
		fields, defaults, err := configSchema.ValidationSchema()
		if err != nil {
			panic(err)
		}
		defaults = updateDefaults(defaults, configBaseDefaults)
		return fields, defaults
	}()

	configImmutableFields = func() []string {
		var names []string
		for name, attr := range configSchema {
			if attr.Immutable {
				names = append(names, name)
			}
		}
		return names
	}()

	configSecretFields = []string{}
)

func updateDefaults(defaults schema.Defaults, updates schema.Defaults) schema.Defaults {
	updated := schema.Defaults{}
	for k, v := range defaults {
		updated[k] = v
	}
	for k, v := range updates {
		// TODO(ericsnow) Delete the item if v is nil?
		updated[k] = v
	}
	return updated
}

func adjustDefaults(cfg *config.Config, defaults map[string]interface{}) (map[string]interface{}, []string) {
	var unset []string
	updated := make(map[string]interface{})
	for k, v := range defaults {
		updated[k] = v
	}

	// The container type would be pulled from cfg if there were more
	// than one possible type for this provider.
	//cType := instance.LXD

	// Set the proper default namespace.
	raw := updated[cfgNamespace]
	if raw == nil || raw.(string) == "" {
		raw = cfg.Name()
		updated[cfgNamespace] = raw
	}

	if val, ok := cfg.UnknownAttrs()[cfgNamespace]; ok && val == "" {
		unset = append(unset, cfgNamespace)
	}

	return updated, unset
}

// TODO(ericsnow) environschema.Fields should have this...
func ensureImmutableFields(oldAttrs, newAttrs map[string]interface{}) error {
	for name, attr := range configSchema {
		if !attr.Immutable {
			continue
		}
		if newAttrs[name] != oldAttrs[name] {
			return errors.Errorf("%s: cannot change from %v to %v", name, oldAttrs[name], newAttrs[name])
		}
	}
	return nil
}

type environConfig struct {
	*config.Config
	attrs map[string]interface{}
}

// newConfig builds a new environConfig from the provided Config and
// returns it.
func newConfig(cfg *config.Config) *environConfig {
	return &environConfig{
		Config: cfg,
		attrs:  cfg.UnknownAttrs(),
	}
}

// newValidConfig builds a new environConfig from the provided Config
// and returns it. This includes applying the provided defaults
// values, if any. The resulting config values are validated.
func newValidConfig(cfg *config.Config, defaults map[string]interface{}) (*environConfig, error) {
	// Any auth credentials handling should happen first...

	// Ensure that the provided config is valid.
	if err := config.Validate(cfg, nil); err != nil {
		return nil, errors.Trace(err)
	}

	// Apply the defaults and coerce/validate the custom config attrs.
	fixedDefaults, unset := adjustDefaults(cfg, defaults)
	cfg, err := cfg.Remove(unset)
	if err != nil {
		return nil, errors.Trace(err)
	}
	validated, err := cfg.ValidateUnknownAttrs(configFields, fixedDefaults)
	if err != nil {
		return nil, errors.Trace(err)
	}
	validCfg, err := cfg.Apply(validated)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Build the config.
	ecfg := newConfig(validCfg)

	// Update to defaults set via client config.
	clientCfg, err := ecfg.clientConfig()
	if err != nil {
		return nil, errors.Trace(err)
	}
	ecfg, err = ecfg.updateForClientConfig(clientCfg)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Do final (more complex, provider-specific) validation.
	if err := ecfg.validate(); err != nil {
		return nil, errors.Trace(err)
	}

	return ecfg, nil
}

func (c *environConfig) containerType() instance.ContainerType {
	// The container type would be pulled from c.attrs if there were
	// more than one possible type for this provider.
	return instance.LXD
}

func (c *environConfig) namespace() string {
	raw := c.attrs[cfgNamespace]
	return raw.(string)
}

func (c *environConfig) dirname() string {
	// TODO(ericsnow) Put it under one of the juju/paths.*() directories.
	return ""
}

func (c *environConfig) remoteURL() string {
	raw := c.attrs[cfgRemoteURL]
	return raw.(string)
}

func (c *environConfig) clientCert() string {
	raw := c.attrs[cfgClientCert]
	return raw.(string)
}

func (c *environConfig) clientKey() string {
	raw := c.attrs[cfgClientKey]
	return raw.(string)
}

// clientConfig builds a LXD Config based on the env config and returns it.
func (c *environConfig) clientConfig() (lxdclient.Config, error) {
	remote := lxdclient.Remote{
		Name: "juju-remote",
		Host: c.remoteURL(),
	}
	if c.clientCert() != "" {
		certPEM := []byte(c.clientCert())
		keyPEM := []byte(c.clientKey())
		remote.Cert = lxdclient.NewCert(certPEM, keyPEM)
		remote.Cert.Name = fmt.Sprintf("juju cert for env %q", c.Name())
	}

	cfg := lxdclient.Config{
		Namespace: c.namespace(),
		Dirname:   c.dirname(),
		Remote:    remote,
	}
	cfg, err := cfg.WithDefaults()
	if err != nil {
		return cfg, errors.Trace(err)
	}
	return cfg, nil
}

// TODO(ericsnow) Switch to a DI testing approach and eliminiate this var.
var asNonLocal = lxdclient.Config.AsNonLocal

func (c *environConfig) updateForClientConfig(clientCfg lxdclient.Config) (*environConfig, error) {
	nonlocal, err := asNonLocal(clientCfg)
	if err != nil {
		return nil, errors.Trace(err)
	}
	clientCfg = nonlocal

	c.attrs[cfgNamespace] = clientCfg.Namespace

	c.attrs[cfgRemoteURL] = clientCfg.Remote.Host

	var cert lxdclient.Cert
	if clientCfg.Remote.Cert != nil {
		cert = *clientCfg.Remote.Cert
	}
	c.attrs[cfgClientCert] = string(cert.CertPEM)
	c.attrs[cfgClientKey] = string(cert.KeyPEM)

	// Apply the updates.
	cfg, err := c.Config.Apply(c.attrs)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return newConfig(cfg), nil
}

// secret gathers the "secret" config values and returns them.
func (c *environConfig) secret() map[string]string {
	if len(configSecretFields) == 0 {
		return nil
	}

	secretAttrs := make(map[string]string, len(configSecretFields))
	for _, key := range configSecretFields {
		secretAttrs[key] = c.attrs[key].(string)
	}
	return secretAttrs
}

// validate checks more complex LCD-specific config values.
func (c environConfig) validate() error {
	// All fields must be populated, even with just the default.
	// TODO(ericsnow) Shouldn't configSchema support this?
	for field := range configFields {
		if dflt, ok := configDefaults[field]; ok && dflt == "" {
			continue
		}
		if c.attrs[field].(string) == "" {
			return errors.Errorf("%s: must not be empty", field)
		}
	}

	// If cert is provided then key must be (and vice versa).
	if c.clientCert() == "" && c.clientKey() != "" {
		return errors.Errorf("missing %s (got %s value %q)", cfgClientCert, cfgClientKey, c.clientKey())
	}
	if c.clientCert() != "" && c.clientKey() == "" {
		return errors.Errorf("missing %s (got %s value %q)", cfgClientKey, cfgClientCert, c.clientCert())
	}

	// Check sanity of complex provider-specific fields.
	cfg, err := c.clientConfig()
	if err != nil {
		return errors.Trace(err)
	}
	if err := cfg.Validate(); err != nil {
		return errors.Trace(err)
	}

	return nil
}

// update applies changes from the provided config to the env config.
// Changes to any immutable attributes result in an error.
func (c *environConfig) update(cfg *config.Config) error {
	// Validate the updates. newValidConfig does not modify the "known"
	// config attributes so it is safe to call Validate here first.
	if err := config.Validate(cfg, c.Config); err != nil {
		return errors.Trace(err)
	}

	updates, err := newValidConfig(cfg, configDefaults)
	if err != nil {
		return errors.Trace(err)
	}

	// Check that no immutable fields have changed.
	attrs := updates.UnknownAttrs()
	if err := ensureImmutableFields(c.attrs, attrs); err != nil {
		return errors.Trace(err)
	}

	// Apply the updates.
	// TODO(ericsnow) Should updates.Config be set in instead of cfg?
	c.Config = cfg
	c.attrs = cfg.UnknownAttrs()
	return nil
}
<commit_msg>Propagate lxdclient changes to the provider.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build go1.3

package lxd

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/schema"
	"gopkg.in/juju/environschema.v1"

	"github.com/juju/juju/container/lxd/lxdclient"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/instance"
)

// TODO(ericsnow) Support providing cert/key file.

// The LXD-specific config keys.
const (
	cfgNamespace  = "namespace"
	cfgRemoteURL  = "remote-url"
	cfgClientCert = "client-cert"
	cfgClientKey  = "client-key"
)

// TODO(ericsnow) Use configSchema.ExampleYAML (once it is implemented)
// to generate boilerplaceConfig.

// boilerplateConfig will be shown in help output, so please keep it up to
// date when you change environment configuration below.
var boilerplateConfig = `
lxd:
    type: lxd

    # namespace identifies the namespace to associate with containers
    # created by the provider.  It is prepended to the container names.
    # By default the environment's name is used as the namespace.
    #
    # namespace: lxd

    # remote-url is the URL to the LXD API server to use for managing
    # containers, if any. If not specified then the locally running LXD
    # server is used.
    #
    # remote-url:

    # The cert and key the client should use to connect to the remote
    # may also be provided. If not then they are auto-generated.
    #
    # client-cert:
    # client-key:

`[1:]

// configSchema defines the schema for the configuration attributes
// defined by the LXD provider.
var configSchema = environschema.Fields{
	cfgNamespace: {
		Description: `Identifies the namespace to associate with containers created by the provider.  It is prepended to the container names.  By default the environment's name is used as the namespace.`,
		Type:        environschema.Tstring,
		Immutable:   true,
	},
	cfgRemoteURL: {
		Description: `Identifies the LXD API server to use for managing containers, if any.`,
		Type:        environschema.Tstring,
		Immutable:   true,
	},
	cfgClientKey: {
		Description: `The client key used for connecting to a LXD host machine.`,
		Type:        environschema.Tstring,
		Immutable:   true,
	},
	cfgClientCert: {
		Description: `The client cert used for connecting to a LXD host machine.`,
		Type:        environschema.Tstring,
		Immutable:   true,
	},
}

var (
	// TODO(ericsnow) Extract the defaults from configSchema as soon as
	// (or if) environschema.Attr supports defaults.

	configBaseDefaults = schema.Defaults{
		cfgNamespace:  "",
		cfgRemoteURL:  "",
		cfgClientCert: "",
		cfgClientKey:  "",
	}

	configFields, configDefaults = func() (schema.Fields, schema.Defaults) {
		fields, defaults, err := configSchema.ValidationSchema()
		if err != nil {
			panic(err)
		}
		defaults = updateDefaults(defaults, configBaseDefaults)
		return fields, defaults
	}()

	configImmutableFields = func() []string {
		var names []string
		for name, attr := range configSchema {
			if attr.Immutable {
				names = append(names, name)
			}
		}
		return names
	}()

	configSecretFields = []string{}
)

func updateDefaults(defaults schema.Defaults, updates schema.Defaults) schema.Defaults {
	updated := schema.Defaults{}
	for k, v := range defaults {
		updated[k] = v
	}
	for k, v := range updates {
		// TODO(ericsnow) Delete the item if v is nil?
		updated[k] = v
	}
	return updated
}

func adjustDefaults(cfg *config.Config, defaults map[string]interface{}) (map[string]interface{}, []string) {
	var unset []string
	updated := make(map[string]interface{})
	for k, v := range defaults {
		updated[k] = v
	}

	// The container type would be pulled from cfg if there were more
	// than one possible type for this provider.
	//cType := instance.LXD

	// Set the proper default namespace.
	raw := updated[cfgNamespace]
	if raw == nil || raw.(string) == "" {
		raw = cfg.Name()
		updated[cfgNamespace] = raw
	}

	if val, ok := cfg.UnknownAttrs()[cfgNamespace]; ok && val == "" {
		unset = append(unset, cfgNamespace)
	}

	return updated, unset
}

// TODO(ericsnow) environschema.Fields should have this...
func ensureImmutableFields(oldAttrs, newAttrs map[string]interface{}) error {
	for name, attr := range configSchema {
		if !attr.Immutable {
			continue
		}
		if newAttrs[name] != oldAttrs[name] {
			return errors.Errorf("%s: cannot change from %v to %v", name, oldAttrs[name], newAttrs[name])
		}
	}
	return nil
}

type environConfig struct {
	*config.Config
	attrs map[string]interface{}
}

// newConfig builds a new environConfig from the provided Config and
// returns it.
func newConfig(cfg *config.Config) *environConfig {
	return &environConfig{
		Config: cfg,
		attrs:  cfg.UnknownAttrs(),
	}
}

// newValidConfig builds a new environConfig from the provided Config
// and returns it. This includes applying the provided defaults
// values, if any. The resulting config values are validated.
func newValidConfig(cfg *config.Config, defaults map[string]interface{}) (*environConfig, error) {
	// Any auth credentials handling should happen first...

	// Ensure that the provided config is valid.
	if err := config.Validate(cfg, nil); err != nil {
		return nil, errors.Trace(err)
	}

	// Apply the defaults and coerce/validate the custom config attrs.
	fixedDefaults, unset := adjustDefaults(cfg, defaults)
	cfg, err := cfg.Remove(unset)
	if err != nil {
		return nil, errors.Trace(err)
	}
	validated, err := cfg.ValidateUnknownAttrs(configFields, fixedDefaults)
	if err != nil {
		return nil, errors.Trace(err)
	}
	validCfg, err := cfg.Apply(validated)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Build the config.
	ecfg := newConfig(validCfg)

	// Update to defaults set via client config.
	clientCfg, err := ecfg.clientConfig()
	if err != nil {
		return nil, errors.Trace(err)
	}
	ecfg, err = ecfg.updateForClientConfig(clientCfg)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Do final (more complex, provider-specific) validation.
	if err := ecfg.validate(); err != nil {
		return nil, errors.Trace(err)
	}

	return ecfg, nil
}

func (c *environConfig) containerType() instance.ContainerType {
	// The container type would be pulled from c.attrs if there were
	// more than one possible type for this provider.
	return instance.LXD
}

func (c *environConfig) namespace() string {
	raw := c.attrs[cfgNamespace]
	return raw.(string)
}

func (c *environConfig) dirname() string {
	// TODO(ericsnow) Put it under one of the juju/paths.*() directories.
	return ""
}

func (c *environConfig) remoteURL() string {
	raw := c.attrs[cfgRemoteURL]
	return raw.(string)
}

func (c *environConfig) clientCert() string {
	raw := c.attrs[cfgClientCert]
	return raw.(string)
}

func (c *environConfig) clientKey() string {
	raw := c.attrs[cfgClientKey]
	return raw.(string)
}

// clientConfig builds a LXD Config based on the env config and returns it.
func (c *environConfig) clientConfig() (lxdclient.Config, error) {
	remote := lxdclient.Remote{
		Name: "juju-remote",
		Host: c.remoteURL(),
	}
	if c.clientCert() != "" {
		certPEM := []byte(c.clientCert())
		keyPEM := []byte(c.clientKey())
		cert := lxdclient.NewCert(certPEM, keyPEM)
		cert.Name = fmt.Sprintf("juju cert for env %q", c.Name())
		remote.Cert = &cert
	}

	cfg := lxdclient.Config{
		Namespace: c.namespace(),
		Dirname:   c.dirname(),
		Remote:    remote,
	}
	cfg, err := cfg.WithDefaults()
	if err != nil {
		return cfg, errors.Trace(err)
	}
	return cfg, nil
}

// TODO(ericsnow) Switch to a DI testing approach and eliminiate this var.
var asNonLocal = lxdclient.Config.UsingTCPRemote

func (c *environConfig) updateForClientConfig(clientCfg lxdclient.Config) (*environConfig, error) {
	nonlocal, err := asNonLocal(clientCfg)
	if err != nil {
		return nil, errors.Trace(err)
	}
	clientCfg = nonlocal

	c.attrs[cfgNamespace] = clientCfg.Namespace

	c.attrs[cfgRemoteURL] = clientCfg.Remote.Host

	var cert lxdclient.Cert
	if clientCfg.Remote.Cert != nil {
		cert = *clientCfg.Remote.Cert
	}
	c.attrs[cfgClientCert] = string(cert.CertPEM)
	c.attrs[cfgClientKey] = string(cert.KeyPEM)

	// Apply the updates.
	cfg, err := c.Config.Apply(c.attrs)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return newConfig(cfg), nil
}

// secret gathers the "secret" config values and returns them.
func (c *environConfig) secret() map[string]string {
	if len(configSecretFields) == 0 {
		return nil
	}

	secretAttrs := make(map[string]string, len(configSecretFields))
	for _, key := range configSecretFields {
		secretAttrs[key] = c.attrs[key].(string)
	}
	return secretAttrs
}

// validate checks more complex LCD-specific config values.
func (c environConfig) validate() error {
	// All fields must be populated, even with just the default.
	// TODO(ericsnow) Shouldn't configSchema support this?
	for field := range configFields {
		if dflt, ok := configDefaults[field]; ok && dflt == "" {
			continue
		}
		if c.attrs[field].(string) == "" {
			return errors.Errorf("%s: must not be empty", field)
		}
	}

	// If cert is provided then key must be (and vice versa).
	if c.clientCert() == "" && c.clientKey() != "" {
		return errors.Errorf("missing %s (got %s value %q)", cfgClientCert, cfgClientKey, c.clientKey())
	}
	if c.clientCert() != "" && c.clientKey() == "" {
		return errors.Errorf("missing %s (got %s value %q)", cfgClientKey, cfgClientCert, c.clientCert())
	}

	// Check sanity of complex provider-specific fields.
	cfg, err := c.clientConfig()
	if err != nil {
		return errors.Trace(err)
	}
	if err := cfg.Validate(); err != nil {
		return errors.Trace(err)
	}

	return nil
}

// update applies changes from the provided config to the env config.
// Changes to any immutable attributes result in an error.
func (c *environConfig) update(cfg *config.Config) error {
	// Validate the updates. newValidConfig does not modify the "known"
	// config attributes so it is safe to call Validate here first.
	if err := config.Validate(cfg, c.Config); err != nil {
		return errors.Trace(err)
	}

	updates, err := newValidConfig(cfg, configDefaults)
	if err != nil {
		return errors.Trace(err)
	}

	// Check that no immutable fields have changed.
	attrs := updates.UnknownAttrs()
	if err := ensureImmutableFields(c.attrs, attrs); err != nil {
		return errors.Trace(err)
	}

	// Apply the updates.
	// TODO(ericsnow) Should updates.Config be set in instead of cfg?
	c.Config = cfg
	c.attrs = cfg.UnknownAttrs()
	return nil
}
<|endoftext|>