<commit_before>package ec2

//----------------------------------------------------------------------------
// Package factored import statement:
//----------------------------------------------------------------------------

import (

	// Stdlib:
	"encoding/base64"
	"fmt"
	"log"
	"strings"

	// Community:
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
)

//----------------------------------------------------------------------------
// Typedefs:
//----------------------------------------------------------------------------

// Data contains variables used by EC2 API.
type Data struct {
	Region             string
	SubnetIDs          string
	ImageID            string
	KeyPair            string
	InstanceType       string
	Hostname           string
	ElasticIP          string
	VpcCidrBlock       string
	VpcID              string
	VpcNameTag         string
	InternalSubnetCidr string
	ExternalSubnetCidr string
	InternalSubnetID   string
	ExternalSubnetID   string
	InternetGatewayID  string
	AllocationID       string
}

//--------------------------------------------------------------------------
// func: Setup
//--------------------------------------------------------------------------

// Setup an EC2 VPC and all the related components.
func (d *Data) Setup() error {

	// Connect and authenticate to the API endpoint:
	log.Printf("[setup-ec2] INFO Connecting to %s\n", d.Region)
	svc := ec2.New(session.New(&aws.Config{Region: aws.String(d.Region)}))

	// Create the VPC:
	if err := d.createVpc(*svc); err != nil {
		return err
	}

	// Create external and internal subnets:
	if err := d.createSubnets(*svc); err != nil {
		return err
	}

	// Create the internet gateway:
	if err := d.createInternetGateway(*svc); err != nil {
		return err
	}

	// Allocate a new EIP:
	if err := d.allocateAddress(*svc); err != nil {
		return err
	}

	// Return on success:
	return nil
}

//-------------------------------------------------------------------------
// func: createVpc
//-------------------------------------------------------------------------

func (d *Data) createVpc(svc ec2.EC2) error {

	// Forge the VPC request:
	prmsVpc := &ec2.CreateVpcInput{
		CidrBlock:       aws.String(d.VpcCidrBlock),
		DryRun:          aws.Bool(false),
		InstanceTenancy: aws.String("default"),
	}

	// Send the VPC request:
	rspVpc, err := svc.CreateVpc(prmsVpc)
	if err != nil {
		return err
	}

	// Store the VPC ID:
	d.VpcID = *rspVpc.Vpc.VpcId
	log.Printf("[setup-ec2] INFO New VPC %s\n", d.VpcID)

	// Tag the VPC:
	if err = tag(d.VpcID, "Name", d.VpcNameTag, svc); err != nil {
		return err
	}

	return nil
}

//-------------------------------------------------------------------------
// func: createSubnets
//-------------------------------------------------------------------------

func (d *Data) createSubnets(svc ec2.EC2) error {

	// Map to iterate:
	nets := map[string]map[string]string{
		"internal": map[string]string{"SubnetCidr": d.InternalSubnetCidr, "SubnetID": ""},
		"external": map[string]string{"SubnetCidr": d.ExternalSubnetCidr, "SubnetID": ""},
	}

	// For each subnet:
	for k, v := range nets {

		// Forge the subnet request:
		params := &ec2.CreateSubnetInput{
			CidrBlock: aws.String(v["SubnetCidr"]),
			VpcId:     aws.String(d.VpcID),
			DryRun:    aws.Bool(false),
		}

		// Send the subnet request:
		resp, err := svc.CreateSubnet(params)
		if err != nil {
			return err
		}

		// Locally store the subnet ID:
		v["SubnetID"] = *resp.Subnet.SubnetId
		log.Printf("[setup-ec2] INFO New %s subnet %s\n", k, v["SubnetID"])

		// Tag the subnet:
		if err = tag(v["SubnetID"], "Name", k, svc); err != nil {
			return err
		}
	}

	// Store subnet IDs:
	d.InternalSubnetID = nets["internal"]["SubnetID"]
	d.ExternalSubnetID = nets["external"]["SubnetID"]

	return nil
}

//-------------------------------------------------------------------------
// func: createInternetGateway
//-------------------------------------------------------------------------

func (d *Data) createInternetGateway(svc ec2.EC2) error {

	// Forge the internet gateway request:
	params := &ec2.CreateInternetGatewayInput{
		DryRun: aws.Bool(false),
	}

	// Send the internet gateway request:
	resp, err := svc.CreateInternetGateway(params)
	if err != nil {
		return err
	}

	// Store the internet gateway ID:
	d.InternetGatewayID = *resp.InternetGateway.InternetGatewayId
	log.Printf("[setup-ec2] INFO New internet gateway %s\n", d.InternetGatewayID)

	return nil
}

//-------------------------------------------------------------------------
// func: allocateAddress
//-------------------------------------------------------------------------

func (d *Data) allocateAddress(svc ec2.EC2) error {

	// Forge the allocation request:
	params := &ec2.AllocateAddressInput{
		Domain: aws.String("vpc"),
		DryRun: aws.Bool(false),
	}

	// Send the allocation request:
	resp, err := svc.AllocateAddress(params)
	if err != nil {
		return err
	}

	// Store the EIP ID:
	d.AllocationID = *resp.AllocationId
	log.Printf("[setup-ec2] INFO New elastic IP %s\n", d.AllocationID)

	return nil
}

//-------------------------------------------------------------------------
// func: tag
//-------------------------------------------------------------------------

func tag(resource, key, value string, svc ec2.EC2) error {

	// Forge the tag request:
	params := &ec2.CreateTagsInput{
		Resources: []*string{
			aws.String(resource),
		},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String(key),
				Value: aws.String(value),
			},
		},
		DryRun: aws.Bool(false),
	}

	// Send the tag request:
	_, err := svc.CreateTags(params)
	if err != nil {
		return err
	}

	return nil
}

//--------------------------------------------------------------------------
// func: Run
//--------------------------------------------------------------------------

// Run uses EC2 API to launch a new instance.
func (d *Data) Run(udata []byte) error {

	// Connect and authenticate to the API endpoint:
	svc := ec2.New(session.New(&aws.Config{Region: aws.String(d.Region)}))

	// Forge the network interfaces:
	var networkInterfaces []*ec2.InstanceNetworkInterfaceSpecification
	subnetIDs := strings.Split(d.SubnetIDs, ",")

	for i := 0; i < len(subnetIDs); i++ {

		// Forge the security group ids:
		var securityGroupIds []*string
		for _, gid := range strings.Split(subnetIDs[i], ":")[1:] {
			securityGroupIds = append(securityGroupIds, aws.String(gid))
		}

		iface := ec2.InstanceNetworkInterfaceSpecification{
			DeleteOnTermination: aws.Bool(true),
			DeviceIndex:         aws.Int64(int64(i)),
			Groups:              securityGroupIds,
			SubnetId:            aws.String(strings.Split(subnetIDs[i], ":")[0]),
		}

		networkInterfaces = append(networkInterfaces, &iface)
	}

	// Send the request:
	runResult, err := svc.RunInstances(&ec2.RunInstancesInput{
		ImageId:           aws.String(d.ImageID),
		MinCount:          aws.Int64(1),
		MaxCount:          aws.Int64(1),
		KeyName:           aws.String(d.KeyPair),
		InstanceType:      aws.String(d.InstanceType),
		NetworkInterfaces: networkInterfaces,
		UserData:          aws.String(base64.StdEncoding.EncodeToString([]byte(udata))),
	})

	if err != nil {
		return err
	}

	// Pretty-print the response data:
	fmt.Println("Created instance", *runResult.Instances[0].InstanceId)

	// Add tags to the created instance:
	_, err = svc.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{runResult.Instances[0].InstanceId},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("Name"),
				Value: aws.String(d.Hostname),
			},
		},
	})

	if err != nil {
		return err
	}

	// Allocate an elastic IP address:
	if d.ElasticIP == "true" {

		params := &ec2.AllocateAddressInput{
			Domain: aws.String("vpc"),
			DryRun: aws.Bool(false),
		}

		// Send the request:
		resp, err := svc.AllocateAddress(params)
		if err != nil {
			return err
		}

		// Pretty-print the response data:
		fmt.Println(resp)
	}

	// Return on success:
	return nil
}
<commit_msg>New func createNatGateway()<commit_after>package ec2

//----------------------------------------------------------------------------
// Package factored import statement:
//----------------------------------------------------------------------------

import (

	// Stdlib:
	"encoding/base64"
	"fmt"
	"log"
	"strings"

	// Community:
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
)

//----------------------------------------------------------------------------
// Typedefs:
//----------------------------------------------------------------------------

// Data contains variables used by EC2 API.
type Data struct {
	Region             string
	SubnetIDs          string
	ImageID            string
	KeyPair            string
	InstanceType       string
	Hostname           string
	ElasticIP          string
	VpcCidrBlock       string
	VpcID              string
	VpcNameTag         string
	InternalSubnetCidr string
	ExternalSubnetCidr string
	InternalSubnetID   string
	ExternalSubnetID   string
	InternetGatewayID  string
	AllocationID       string
	NatGatewayID       string
}

//--------------------------------------------------------------------------
// func: Setup
//--------------------------------------------------------------------------

// Setup an EC2 VPC and all the related components.
func (d *Data) Setup() error {

	// Connect and authenticate to the API endpoint:
	log.Printf("[setup-ec2] INFO Connecting to %s\n", d.Region)
	svc := ec2.New(session.New(&aws.Config{Region: aws.String(d.Region)}))

	// Create the VPC:
	if err := d.createVpc(*svc); err != nil {
		return err
	}

	// Create external and internal subnets:
	if err := d.createSubnets(*svc); err != nil {
		return err
	}

	// Create the internet gateway:
	if err := d.createInternetGateway(*svc); err != nil {
		return err
	}

	// Allocate a new EIP:
	if err := d.allocateAddress(*svc); err != nil {
		return err
	}

	// Create a NAT gateway:
	if err := d.createNatGateway(*svc); err != nil {
		return err
	}

	// Return on success:
	return nil
}

//-------------------------------------------------------------------------
// func: createVpc
//-------------------------------------------------------------------------

func (d *Data) createVpc(svc ec2.EC2) error {

	// Forge the VPC request:
	prmsVpc := &ec2.CreateVpcInput{
		CidrBlock:       aws.String(d.VpcCidrBlock),
		DryRun:          aws.Bool(false),
		InstanceTenancy: aws.String("default"),
	}

	// Send the VPC request:
	rspVpc, err := svc.CreateVpc(prmsVpc)
	if err != nil {
		return err
	}

	// Store the VPC ID:
	d.VpcID = *rspVpc.Vpc.VpcId
	log.Printf("[setup-ec2] INFO New VPC %s\n", d.VpcID)

	// Tag the VPC:
	if err = tag(d.VpcID, "Name", d.VpcNameTag, svc); err != nil {
		return err
	}

	return nil
}

//-------------------------------------------------------------------------
// func: createSubnets
//-------------------------------------------------------------------------

func (d *Data) createSubnets(svc ec2.EC2) error {

	// Map to iterate:
	nets := map[string]map[string]string{
		"internal": map[string]string{"SubnetCidr": d.InternalSubnetCidr, "SubnetID": ""},
		"external": map[string]string{"SubnetCidr": d.ExternalSubnetCidr, "SubnetID": ""},
	}

	// For each subnet:
	for k, v := range nets {

		// Forge the subnet request:
		params := &ec2.CreateSubnetInput{
			CidrBlock: aws.String(v["SubnetCidr"]),
			VpcId:     aws.String(d.VpcID),
			DryRun:    aws.Bool(false),
		}

		// Send the subnet request:
		resp, err := svc.CreateSubnet(params)
		if err != nil {
			return err
		}

		// Locally store the subnet ID:
		v["SubnetID"] = *resp.Subnet.SubnetId
		log.Printf("[setup-ec2] INFO New %s subnet %s\n", k, v["SubnetID"])

		// Tag the subnet:
		if err = tag(v["SubnetID"], "Name", k, svc); err != nil {
			return err
		}
	}

	// Store subnet IDs:
	d.InternalSubnetID = nets["internal"]["SubnetID"]
	d.ExternalSubnetID = nets["external"]["SubnetID"]

	return nil
}

//-------------------------------------------------------------------------
// func: createInternetGateway
//-------------------------------------------------------------------------

func (d *Data) createInternetGateway(svc ec2.EC2) error {

	// Forge the internet gateway request:
	params := &ec2.CreateInternetGatewayInput{
		DryRun: aws.Bool(false),
	}

	// Send the internet gateway request:
	resp, err := svc.CreateInternetGateway(params)
	if err != nil {
		return err
	}

	// Store the internet gateway ID:
	d.InternetGatewayID = *resp.InternetGateway.InternetGatewayId
	log.Printf("[setup-ec2] INFO New internet gateway %s\n", d.InternetGatewayID)

	return nil
}

//-------------------------------------------------------------------------
// func: allocateAddress
//-------------------------------------------------------------------------

func (d *Data) allocateAddress(svc ec2.EC2) error {

	// Forge the allocation request:
	params := &ec2.AllocateAddressInput{
		Domain: aws.String("vpc"),
		DryRun: aws.Bool(false),
	}

	// Send the allocation request:
	resp, err := svc.AllocateAddress(params)
	if err != nil {
		return err
	}

	// Store the EIP ID:
	d.AllocationID = *resp.AllocationId
	log.Printf("[setup-ec2] INFO New elastic IP %s\n", d.AllocationID)

	return nil
}

//-------------------------------------------------------------------------
// func: createNatGateway
//-------------------------------------------------------------------------

func (d *Data) createNatGateway(svc ec2.EC2) error {

	// Forge the NAT gateway request:
	params := &ec2.CreateNatGatewayInput{
		AllocationId: aws.String(d.AllocationID),
		SubnetId:     aws.String(d.ExternalSubnetID),
		ClientToken:  aws.String("kato"),
	}

	// Send the NAT gateway request:
	resp, err := svc.CreateNatGateway(params)
	if err != nil {
		return err
	}

	// Store the NAT gateway ID:
	d.NatGatewayID = *resp.NatGateway.NatGatewayId
	log.Printf("[setup-ec2] INFO New NAT gateway %s\n", d.NatGatewayID)

	return nil
}

//-------------------------------------------------------------------------
// func: tag
//-------------------------------------------------------------------------

func tag(resource, key, value string, svc ec2.EC2) error {

	// Forge the tag request:
	params := &ec2.CreateTagsInput{
		Resources: []*string{
			aws.String(resource),
		},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String(key),
				Value: aws.String(value),
			},
		},
		DryRun: aws.Bool(false),
	}

	// Send the tag request:
	_, err := svc.CreateTags(params)
	if err != nil {
		return err
	}

	return nil
}

//--------------------------------------------------------------------------
// func: Run
//--------------------------------------------------------------------------

// Run uses EC2 API to launch a new instance.
func (d *Data) Run(udata []byte) error {

	// Connect and authenticate to the API endpoint:
	svc := ec2.New(session.New(&aws.Config{Region: aws.String(d.Region)}))

	// Forge the network interfaces:
	var networkInterfaces []*ec2.InstanceNetworkInterfaceSpecification
	subnetIDs := strings.Split(d.SubnetIDs, ",")

	for i := 0; i < len(subnetIDs); i++ {

		// Forge the security group ids:
		var securityGroupIds []*string
		for _, gid := range strings.Split(subnetIDs[i], ":")[1:] {
			securityGroupIds = append(securityGroupIds, aws.String(gid))
		}

		iface := ec2.InstanceNetworkInterfaceSpecification{
			DeleteOnTermination: aws.Bool(true),
			DeviceIndex:         aws.Int64(int64(i)),
			Groups:              securityGroupIds,
			SubnetId:            aws.String(strings.Split(subnetIDs[i], ":")[0]),
		}

		networkInterfaces = append(networkInterfaces, &iface)
	}

	// Send the request:
	runResult, err := svc.RunInstances(&ec2.RunInstancesInput{
		ImageId:           aws.String(d.ImageID),
		MinCount:          aws.Int64(1),
		MaxCount:          aws.Int64(1),
		KeyName:           aws.String(d.KeyPair),
		InstanceType:      aws.String(d.InstanceType),
		NetworkInterfaces: networkInterfaces,
		UserData:          aws.String(base64.StdEncoding.EncodeToString([]byte(udata))),
	})

	if err != nil {
		return err
	}

	// Pretty-print the response data:
	fmt.Println("Created instance", *runResult.Instances[0].InstanceId)

	// Add tags to the created instance:
	_, err = svc.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{runResult.Instances[0].InstanceId},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("Name"),
				Value: aws.String(d.Hostname),
			},
		},
	})

	if err != nil {
		return err
	}

	// Allocate an elastic IP address:
	if d.ElasticIP == "true" {

		params := &ec2.AllocateAddressInput{
			Domain: aws.String("vpc"),
			DryRun: aws.Bool(false),
		}

		// Send the request:
		resp, err := svc.AllocateAddress(params)
		if err != nil {
			return err
		}

		// Pretty-print the response data:
		fmt.Println(resp)
	}

	// Return on success:
	return nil
}
<|endoftext|>


<commit_before>package profile

import (
	"github.com/hackform/governor"
	"github.com/hackform/governor/service/cache"
	"github.com/hackform/governor/service/db"
	"github.com/hackform/governor/service/objstore"
	"github.com/hackform/governor/service/profile/model"
	"github.com/hackform/governor/service/user/gate"
	"github.com/hackform/governor/util/uid"
	"github.com/labstack/echo"
	"github.com/sirupsen/logrus"
	"io"
	"mime"
	"net/http"
	"time"
)

type (
	reqProfileGetID struct {
		Userid string `json:"userid"`
	}

	reqProfileModel struct {
		Userid string `json:"-"`
		Email  string `json:"contact_email"`
		Bio    string `json:"bio"`
	}

	reqProfileImage struct {
		Userid    string `json:"-"`
		Image     io.Reader
		ImageType string
		Key       string
	}

	resProfileUpdate struct {
		Userid string `json:"userid"`
	}

	resProfileModel struct {
		Email string `json:"contact_email"`
		Bio   string `json:"bio"`
		Image string `json:"image"`
	}

	resProfileImageChange struct {
		Key string `json:"key"`
	}
)

func (r *reqProfileGetID) valid() *governor.Error {
	if err := hasUserid(r.Userid); err != nil {
		return err
	}
	return nil
}

func (r *reqProfileModel) valid() *governor.Error {
	if err := hasUserid(r.Userid); err != nil {
		return err
	}
	if err := validEmail(r.Email); err != nil {
		return err
	}
	if err := validBio(r.Email); err != nil {
		return err
	}
	return nil
}

func (r *reqProfileImage) valid() *governor.Error {
	if err := hasUserid(r.Userid); err != nil {
		return err
	}
	if err := validImage(r.Image); err != nil {
		return err
	}
	if err := validImageType(r.ImageType); err != nil {
		return err
	}
	if err := validKey(r.Key); err != nil {
		return err
	}
	return nil
}

const (
	moduleID    = "profile"
	imageBucket = "profile-image"
)

type (
	// Profile is a service for storing user profile information
	Profile struct {
		db    *db.Database
		cache *cache.Cache
		obj   *objstore.Bucket
		gate  *gate.Gate
	}
)

// New creates a new Profile service
func New(conf governor.Config, l *logrus.Logger, db *db.Database, ch *cache.Cache, obj *objstore.Objstore) *Profile {
	ca := conf.Conf().GetStringMapString("userauth")

	b, err := obj.GetBucketDefLoc(imageBucket)
	if err != nil {
		l.Errorf("failed to get bucket: %s\n", err.Error())
	}

	l.Info("initialized profile service")

	return &Profile{
		db:    db,
		cache: ch,
		obj:   b,
		gate:  gate.New(ca["secret"], ca["issuer"]),
	}
}

// Mount is a collection of routes for accessing and modifying profile data
func (p *Profile) Mount(conf governor.Config, r *echo.Group, l *logrus.Logger) error {
	db := p.db.DB()
	ch := p.cache.Cache()

	r.POST("/:id", func(c echo.Context) error {
		rprofile := &reqProfileModel{}
		if err := c.Bind(rprofile); err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		rprofile.Userid = c.Param("id")
		if err := rprofile.valid(); err != nil {
			return err
		}

		m := &profilemodel.Model{
			Email: rprofile.Email,
			Bio:   rprofile.Bio,
		}

		if err := m.SetIDB64(rprofile.Userid); err != nil {
			err.SetErrorUser()
			return err
		}

		if err := m.Insert(db); err != nil {
			if err.Code() == 3 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}

		userid, _ := m.IDBase64()

		return c.JSON(http.StatusCreated, &resProfileUpdate{
			Userid: userid,
		})
	}, p.gate.Owner("id"))

	r.PUT("/:id", func(c echo.Context) error {
		rprofile := &reqProfileModel{}
		if err := c.Bind(rprofile); err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		rprofile.Userid = c.Param("id")
		if err := rprofile.valid(); err != nil {
			return err
		}

		m, err := profilemodel.GetByIDB64(db, rprofile.Userid)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}

		m.Email = rprofile.Email
		m.Bio = rprofile.Bio

		if err := m.Update(db); err != nil {
			err.AddTrace(moduleID)
			return err
		}

		return c.NoContent(http.StatusNoContent)
	}, p.gate.Owner("id"))

	r.PUT("/:id/image", func(c echo.Context) error {
		userid := c.Get("userid").(string)

		// create a new key for the session
		key, err := uid.NewU(0, 16)
		if err != nil {
			err.AddTrace(moduleID)
			return err
		}
		sessionKey := key.Base64()

		if err := ch.Set(sessionKey, userid, time.Second*120).Err(); err != nil {
			return governor.NewError(moduleID, err.Error(), 0, http.StatusInternalServerError)
		}

		return c.JSON(http.StatusOK, &resProfileImageChange{
			Key: sessionKey,
		})
	}, p.gate.Owner("id"))

	r.POST("/:id/image/:key", func(c echo.Context) error {
		file, err := c.FormFile("image")
		if err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		mediaType, _, err := mime.ParseMediaType(file.Header.Get("Content-Type"))
		if err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		image, err := file.Open()
		if err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		defer func(c io.Closer) {
			err := c.Close()
			if err != nil {
				gerr := governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
				l.WithFields(logrus.Fields{
					"origin": gerr.Origin(),
					"source": gerr.Source(),
					"code":   gerr.Code(),
				}).Error(gerr.Message())
			}
		}(image)
		rprofile := &reqProfileImage{
			Userid:    c.Param("id"),
			Image:     image,
			ImageType: mediaType,
			Key:       c.Param("key"),
		}
		if err := rprofile.valid(); err != nil {
			return err
		}

		userid, err := ch.Get(rprofile.Key).Result()
		if err != nil {
			return governor.NewErrorUser(moduleID, "key not found: "+err.Error(), 0, http.StatusBadRequest)
		}

		if rprofile.Userid != userid {
			return governor.NewErrorUser(moduleID, "improper profile id", 0, http.StatusBadRequest)
		}

		if err := ch.Del(rprofile.Key).Err(); err != nil {
			return governor.NewError(moduleID, err.Error(), 0, http.StatusInternalServerError)
		}

		if err := p.obj.Put(rprofile.Userid+"-profile", rprofile.ImageType, image); err != nil {
			err.AddTrace(moduleID)
			return err
		}

		return c.NoContent(http.StatusNoContent)
	})

	r.DELETE("/:id", func(c echo.Context) error {
		rprofile := &reqProfileGetID{
			Userid: c.Param("id"),
		}
		if err := rprofile.valid(); err != nil {
			return err
		}

		m, err := profilemodel.GetByIDB64(db, rprofile.Userid)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}

		if err := m.Delete(db); err != nil {
			err.AddTrace(moduleID)
			return err
		}

		return c.NoContent(http.StatusNoContent)
	}, p.gate.OwnerOrAdmin("id"))

	r.GET("/:id", func(c echo.Context) error {
		rprofile := &reqProfileGetID{
			Userid: c.Param("id"),
		}
		if err := rprofile.valid(); err != nil {
			return err
		}

		m, err := profilemodel.GetByIDB64(db, rprofile.Userid)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}

		return c.JSON(http.StatusOK, &resProfileModel{
			Email: m.Email,
			Bio:   m.Bio,
			Image: m.Image,
		})
	})

	r.GET("/:id/image", func(c echo.Context) error {
		rprofile := &reqProfileGetID{
			Userid: c.Param("id"),
		}
		if err := rprofile.valid(); err != nil {
			return err
		}

		obj, objinfo, err := p.obj.Get(rprofile.Userid + "-profile")
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}
		return c.Stream(http.StatusOK, objinfo.ContentType, obj)
	})

	l.Info("mounted profile service")

	return nil
}

// Health is a check for service health
func (p *Profile) Health() *governor.Error {
	return nil
}

// Setup is run on service setup
func (p *Profile) Setup(conf governor.Config, l *logrus.Logger, rsetup governor.ReqSetupPost) *governor.Error {
	if err := profilemodel.Setup(p.db.DB()); err != nil {
		err.AddTrace(moduleID)
		return err
	}
	l.Info("created new profile table")
	return nil
}
<commit_msg>profile image upload only allow one outstanding key<commit_after>package profile

import (
	"github.com/hackform/governor"
	"github.com/hackform/governor/service/cache"
	"github.com/hackform/governor/service/db"
	"github.com/hackform/governor/service/objstore"
	"github.com/hackform/governor/service/profile/model"
	"github.com/hackform/governor/service/user/gate"
	"github.com/hackform/governor/util/uid"
	"github.com/labstack/echo"
	"github.com/sirupsen/logrus"
	"io"
	"mime"
	"net/http"
	"time"
)

type (
	reqProfileGetID struct {
		Userid string `json:"userid"`
	}

	reqProfileModel struct {
		Userid string `json:"-"`
		Email  string `json:"contact_email"`
		Bio    string `json:"bio"`
	}

	reqProfileImage struct {
		Userid    string `json:"-"`
		Image     io.Reader
		ImageType string
		Key       string
	}

	resProfileUpdate struct {
		Userid string `json:"userid"`
	}

	resProfileModel struct {
		Email string `json:"contact_email"`
		Bio   string `json:"bio"`
		Image string `json:"image"`
	}

	resProfileImageChange struct {
		Key string `json:"key"`
	}
)

func (r *reqProfileGetID) valid() *governor.Error {
	if err := hasUserid(r.Userid); err != nil {
		return err
	}
	return nil
}

func (r *reqProfileModel) valid() *governor.Error {
	if err := hasUserid(r.Userid); err != nil {
		return err
	}
	if err := validEmail(r.Email); err != nil {
		return err
	}
	if err := validBio(r.Email); err != nil {
		return err
	}
	return nil
}

func (r *reqProfileImage) valid() *governor.Error {
	if err := hasUserid(r.Userid); err != nil {
		return err
	}
	if err := validImage(r.Image); err != nil {
		return err
	}
	if err := validImageType(r.ImageType); err != nil {
		return err
	}
	if err := validKey(r.Key); err != nil {
		return err
	}
	return nil
}

const (
	moduleID    = "profile"
	imageBucket = "profile-image"
)

type (
	// Profile is a service for storing user profile information
	Profile struct {
		db    *db.Database
		cache *cache.Cache
		obj   *objstore.Bucket
		gate  *gate.Gate
	}
)

// New creates a new Profile service
func New(conf governor.Config, l *logrus.Logger, db *db.Database, ch *cache.Cache, obj *objstore.Objstore) *Profile {
	ca := conf.Conf().GetStringMapString("userauth")

	b, err := obj.GetBucketDefLoc(imageBucket)
	if err != nil {
		l.Errorf("failed to get bucket: %s\n", err.Error())
	}

	l.Info("initialized profile service")

	return &Profile{
		db:    db,
		cache: ch,
		obj:   b,
		gate:  gate.New(ca["secret"], ca["issuer"]),
	}
}

// Mount is a collection of routes for accessing and modifying profile data
func (p *Profile) Mount(conf governor.Config, r *echo.Group, l *logrus.Logger) error {
	db := p.db.DB()
	ch := p.cache.Cache()

	r.POST("/:id", func(c echo.Context) error {
		rprofile := &reqProfileModel{}
		if err := c.Bind(rprofile); err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		rprofile.Userid = c.Param("id")
		if err := rprofile.valid(); err != nil {
			return err
		}

		m := &profilemodel.Model{
			Email: rprofile.Email,
			Bio:   rprofile.Bio,
		}

		if err := m.SetIDB64(rprofile.Userid); err != nil {
			err.SetErrorUser()
			return err
		}

		if err := m.Insert(db); err != nil {
			if err.Code() == 3 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}

		userid, _ := m.IDBase64()

		return c.JSON(http.StatusCreated, &resProfileUpdate{
			Userid: userid,
		})
	}, p.gate.Owner("id"))

	r.PUT("/:id", func(c echo.Context) error {
		rprofile := &reqProfileModel{}
		if err := c.Bind(rprofile); err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		rprofile.Userid = c.Param("id")
		if err := rprofile.valid(); err != nil {
			return err
		}

		m, err := profilemodel.GetByIDB64(db, rprofile.Userid)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}

		m.Email = rprofile.Email
		m.Bio = rprofile.Bio

		if err := m.Update(db); err != nil {
			err.AddTrace(moduleID)
			return err
		}

		return c.NoContent(http.StatusNoContent)
	}, p.gate.Owner("id"))

	r.PUT("/:id/image", func(c echo.Context) error {
		userid := c.Get("userid").(string)

		// create a new key for the session
		key, err := uid.NewU(0, 16)
		if err != nil {
			err.AddTrace(moduleID)
			return err
		}
		sessionKey := key.Base64()

		if err := ch.Set("profileimage-"+userid, sessionKey, time.Second*300).Err(); err != nil {
			return governor.NewError(moduleID, err.Error(), 0, http.StatusInternalServerError)
		}

		return c.JSON(http.StatusOK, &resProfileImageChange{
			Key: sessionKey,
		})
	}, p.gate.Owner("id"))

	r.POST("/:id/image/:key", func(c echo.Context) error {
		file, err := c.FormFile("image")
		if err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		mediaType, _, err := mime.ParseMediaType(file.Header.Get("Content-Type"))
		if err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		image, err := file.Open()
		if err != nil {
			return governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
		}
		defer func(c io.Closer) {
			err := c.Close()
			if err != nil {
				gerr := governor.NewErrorUser(moduleID, err.Error(), 0, http.StatusBadRequest)
				l.WithFields(logrus.Fields{
					"origin": gerr.Origin(),
					"source": gerr.Source(),
					"code":   gerr.Code(),
				}).Error(gerr.Message())
			}
		}(image)
		rprofile := &reqProfileImage{
			Userid:    c.Param("id"),
			Image:     image,
			ImageType: mediaType,
			Key:       c.Param("key"),
		}
		if err := rprofile.valid(); err != nil {
			return err
		}

		key, err := ch.Get("profileimage-" + rprofile.Userid).Result()
		if err != nil {
			return governor.NewErrorUser(moduleID, "profile not found: "+err.Error(), 0, http.StatusBadRequest)
		}

		if rprofile.Key != key {
			return governor.NewErrorUser(moduleID, "improper key", 0, http.StatusBadRequest)
		}

		if err := ch.Del("profileimage-" + rprofile.Userid).Err(); err != nil {
			return governor.NewError(moduleID, err.Error(), 0, http.StatusInternalServerError)
		}

		if err := p.obj.Put(rprofile.Userid+"-profile", rprofile.ImageType, image); err != nil {
			err.AddTrace(moduleID)
			return err
		}

		return c.NoContent(http.StatusNoContent)
	})

	r.DELETE("/:id", func(c echo.Context) error {
		rprofile := &reqProfileGetID{
			Userid: c.Param("id"),
		}
		if err := rprofile.valid(); err != nil {
			return err
		}

		m, err := profilemodel.GetByIDB64(db, rprofile.Userid)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}

		if err := m.Delete(db); err != nil {
			err.AddTrace(moduleID)
			return err
		}

		return c.NoContent(http.StatusNoContent)
	}, p.gate.OwnerOrAdmin("id"))

	r.GET("/:id", func(c echo.Context) error {
		rprofile := &reqProfileGetID{
			Userid: c.Param("id"),
		}
		if err := rprofile.valid(); err != nil {
			return err
		}

		m, err := profilemodel.GetByIDB64(db, rprofile.Userid)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}

		return c.JSON(http.StatusOK, &resProfileModel{
			Email: m.Email,
			Bio:   m.Bio,
			Image: m.Image,
		})
	})

	r.GET("/:id/image", func(c echo.Context) error {
		rprofile := &reqProfileGetID{
			Userid: c.Param("id"),
		}
		if err := rprofile.valid(); err != nil {
			return err
		}

		obj, objinfo, err := p.obj.Get(rprofile.Userid + "-profile")
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleID)
			return err
		}
		return c.Stream(http.StatusOK, objinfo.ContentType, obj)
	})

	l.Info("mounted profile service")

	return nil
}

// Health is a check for service health
func (p *Profile) Health() *governor.Error {
	return nil
}

// Setup is run on service setup
func (p *Profile) Setup(conf governor.Config, l *logrus.Logger, rsetup governor.ReqSetupPost) *governor.Error {
	if err := profilemodel.Setup(p.db.DB()); err != nil {
		err.AddTrace(moduleID)
		return err
	}
	l.Info("created new profile table")
	return nil
}
<|endoftext|>


<commit_before>package jd

import (
	"testing"
)

func TestEqual(t *testing.T) {
	checkEqual(t, `{"a":1}`, `{"a":1}`)
	checkEqual(t, `{"a":1}`, `{"a":1.0}`)
	checkEqual(t, `{"a":[1,2]}`, `{"a":[1,2]}`)
	checkEqual(t, `{"a":"b"}`, `{"a":"b"}`)
}

func TestNotEqual(t *testing.T) {
	checkNotEqual(t, `{"a":1}`, `{"b":1}`)
	checkNotEqual(t, `{"a":[1,2]}`, `{"a":[2,1]}`)
	checkNotEqual(t, `{"a":"b"}`, `{"a":"c"}`)
}
<commit_msg>Rename Test(Not)Equal to TestStruct(Not)Equal.<commit_after>package jd

import (
	"testing"
)

func TestObjectEqual(t *testing.T) {
	checkEqual(t, `{"a":1}`, `{"a":1}`)
	checkEqual(t, `{"a":1}`, `{"a":1.0}`)
	checkEqual(t, `{"a":[1,2]}`, `{"a":[1,2]}`)
	checkEqual(t, `{"a":"b"}`, `{"a":"b"}`)
}

func TestObjectNotEqual(t *testing.T) {
	checkNotEqual(t, `{"a":1}`, `{"b":1}`)
	checkNotEqual(t, `{"a":[1,2]}`, `{"a":[2,1]}`)
	checkNotEqual(t, `{"a":"b"}`, `{"a":"c"}`)
}
<|endoftext|>


<commit_before>package checkawscloudwatchlogs

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs/cloudwatchlogsiface"
	"github.com/jessevdk/go-flags"

	"github.com/mackerelio/checkers"
	"github.com/mackerelio/golib/pluginutil"
	"github.com/natefinch/atomic"
)

type logOpts struct {
	LogGroupName        string `long:"log-group-name" required:"true" value-name:"LOG-GROUP-NAME" description:"Log group name" unquote:"false"`
	LogStreamNamePrefix string `long:"log-stream-name-prefix" value-name:"LOG-STREAM-NAME-PREFIX" description:"Log stream name prefix" unquote:"false"`

	Pattern       string `short:"p" long:"pattern" required:"true" value-name:"PATTERN" description:"Pattern to search for. The value is recognized as the pattern syntax of CloudWatch Logs." unquote:"false"`
	WarningOver   int    `short:"w" long:"warning-over" value-name:"WARNING" description:"Trigger a warning if matched lines is over a number"`
	CriticalOver  int    `short:"c" long:"critical-over" value-name:"CRITICAL" description:"Trigger a critical if matched lines is over a number"`
	StateDir      string `short:"s" long:"state-dir" value-name:"DIR" description:"Dir to keep state files under" unquote:"false"`
	ReturnContent bool   `short:"r" long:"return" description:"Output matched lines"`
	MaxRetries    int    `short:"t" long:"max-retries" value-name:"MAX-RETRIES" description:"Maximum number of retries to call the AWS API"`
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "CloudWatch Logs"
	ckr.Exit()
}

type awsCloudwatchLogsPlugin struct {
	Service   cloudwatchlogsiface.CloudWatchLogsAPI
	StateFile string
	*logOpts
}

func newCloudwatchLogsPlugin(opts *logOpts, args []string) (*awsCloudwatchLogsPlugin, error) {
	var err error
	p := &awsCloudwatchLogsPlugin{logOpts: opts}
	p.Service, err = createService(opts)
	if err != nil {
		return nil, err
	}
	if p.StateDir == "" {
		workdir := pluginutil.PluginWorkDir()
		p.StateDir = filepath.Join(workdir, "check-cloudwatch-logs")
	}
	p.StateFile = getStateFile(p.StateDir, opts.LogGroupName, opts.LogStreamNamePrefix, args)
	return p, nil
}

var stateRe = regexp.MustCompile(`[^-a-zA-Z0-9_.]`)

func getStateFile(stateDir, logGroupName, logStreamNamePrefix string, args []string) string {
	return filepath.Join(
		stateDir,
		fmt.Sprintf(
			"%s-%x.json",
			strings.TrimLeft(stateRe.ReplaceAllString(logGroupName+"_"+logStreamNamePrefix, "_"), "_"),
			md5.Sum([]byte(
				strings.Join(
					[]string{
						os.Getenv("AWS_PROFILE"),
						os.Getenv("AWS_ACCESS_KEY_ID"),
						os.Getenv("AWS_SECRET_ACCESS_KEY"),
						os.Getenv("AWS_REGION"),
						strings.Join(args, " "),
					},
					" ",
				)),
			),
		),
	)
}

func createAWSConfig(opts *logOpts) *aws.Config {
	conf := aws.NewConfig()
	if opts.MaxRetries > 0 {
		return conf.WithMaxRetries(opts.MaxRetries)
	}
	return conf
}

func createService(opts *logOpts) (*cloudwatchlogs.CloudWatchLogs, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}
	return cloudwatchlogs.New(sess, createAWSConfig(opts)), nil
}

type logState struct {
	NextToken *string
	StartTime *int64
}

func (p *awsCloudwatchLogsPlugin) collect(now time.Time) ([]string, error) {
	s, err := p.loadState()
	if err != nil {
		return nil, err
	}
	if s.StartTime == nil || *s.StartTime <= now.Add(-time.Hour).Unix()*1000 {
		s.StartTime = aws.Int64(now.Add(-1*time.Minute).Unix() * 1000)
	}
	var messages []string
	input := &cloudwatchlogs.FilterLogEventsInput{
		StartTime:     s.StartTime,
		LogGroupName:  aws.String(p.LogGroupName),
		NextToken:     s.NextToken,
		FilterPattern: aws.String(p.Pattern),
	}
	if p.LogStreamNamePrefix != "" {
		input.LogStreamNamePrefix = aws.String(p.LogStreamNamePrefix)
	}
	err = p.Service.FilterLogEventsPages(input, func(output *cloudwatchlogs.FilterLogEventsOutput, lastPage bool) bool {
		for _, event := range output.Events {
			messages = append(messages, *event.Message)
			if s.StartTime == nil || *s.StartTime <= *event.Timestamp {
				s.StartTime = aws.Int64(*event.Timestamp + 1)
			}
		}
		s.NextToken = output.NextToken
		if lastPage {
			s.NextToken = nil
		}
		time.Sleep(150 * time.Millisecond)
		return true
	})
	if err != nil {
		return nil, err
	}
	err = p.saveState(s)
	if err != nil {
		return nil, err
	}
	return messages, nil
}

func (p *awsCloudwatchLogsPlugin) loadState() (*logState, error) {
	var s logState
	f, err := os.Open(p.StateFile)
	defer f.Close()
	if err != nil {
		if os.IsNotExist(err) {
			return &s, nil
		}
		return nil, err
	}
	err = json.NewDecoder(f).Decode(&s)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func (p *awsCloudwatchLogsPlugin) saveState(s *logState) error {
	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(s); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(p.StateFile), 0755); err != nil {
		return err
	}
	return atomic.WriteFile(p.StateFile, &buf)
}

func (p *awsCloudwatchLogsPlugin) check(messages []string) *checkers.Checker {
	status := checkers.OK
	msg := fmt.Sprint(len(messages))
	if len(messages) > p.CriticalOver {
		status = checkers.CRITICAL
		msg += " > " + fmt.Sprint(p.CriticalOver)
	} else if len(messages) > p.WarningOver {
		status = checkers.WARNING
		msg += " > " + fmt.Sprint(p.WarningOver)
	}
	msg += " messages for pattern /" + p.Pattern + "/"
	if status != checkers.OK && p.ReturnContent {
		msg += "\n" + strings.Join(messages, "")
	}
	return checkers.NewChecker(status, msg)
}

func (p *awsCloudwatchLogsPlugin) run(now time.Time) *checkers.Checker {
	messages, err := p.collect(now)
	if err != nil {
		return checkers.Unknown(fmt.Sprint(err))
	}
	return p.check(messages)
}

func run(args []string) *checkers.Checker {
	opts := &logOpts{}
	_, err := flags.ParseArgs(opts, args)
	if err != nil {
		os.Exit(1)
	}
	p, err := newCloudwatchLogsPlugin(opts, args)
	if err != nil {
		return checkers.Unknown(fmt.Sprint(err))
	}
	return p.run(time.Now())
}
<commit_msg>save state in iteration<commit_after>package checkawscloudwatchlogs

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs/cloudwatchlogsiface"
	"github.com/jessevdk/go-flags"

	"github.com/mackerelio/checkers"
	"github.com/mackerelio/golib/pluginutil"
	"github.com/natefinch/atomic"
)

type logOpts struct {
	LogGroupName        string `long:"log-group-name" required:"true" value-name:"LOG-GROUP-NAME" description:"Log group name" unquote:"false"`
	LogStreamNamePrefix string `long:"log-stream-name-prefix" value-name:"LOG-STREAM-NAME-PREFIX" description:"Log stream name prefix" unquote:"false"`

	Pattern       string `short:"p" long:"pattern" required:"true" value-name:"PATTERN" description:"Pattern to search for. The value is recognized as the pattern syntax of CloudWatch Logs." unquote:"false"`
	WarningOver   int    `short:"w" long:"warning-over" value-name:"WARNING" description:"Trigger a warning if matched lines is over a number"`
	CriticalOver  int    `short:"c" long:"critical-over" value-name:"CRITICAL" description:"Trigger a critical if matched lines is over a number"`
	StateDir      string `short:"s" long:"state-dir" value-name:"DIR" description:"Dir to keep state files under" unquote:"false"`
	ReturnContent bool   `short:"r" long:"return" description:"Output matched lines"`
	MaxRetries    int    `short:"t" long:"max-retries" value-name:"MAX-RETRIES" description:"Maximum number of retries to call the AWS API"`
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "CloudWatch Logs"
	ckr.Exit()
}

type awsCloudwatchLogsPlugin struct {
	Service   cloudwatchlogsiface.CloudWatchLogsAPI
	StateFile string
	*logOpts
}

func newCloudwatchLogsPlugin(opts *logOpts, args []string) (*awsCloudwatchLogsPlugin, error) {
	var err error
	p := &awsCloudwatchLogsPlugin{logOpts: opts}
	p.Service, err = createService(opts)
	if err != nil {
		return nil, err
	}
	if p.StateDir == "" {
		workdir := pluginutil.PluginWorkDir()
		p.StateDir = filepath.Join(workdir, "check-cloudwatch-logs")
	}
	p.StateFile = getStateFile(p.StateDir, opts.LogGroupName, opts.LogStreamNamePrefix, args)
	return p, nil
}

var stateRe = regexp.MustCompile(`[^-a-zA-Z0-9_.]`)

func getStateFile(stateDir, logGroupName, logStreamNamePrefix string, args []string) string {
	return filepath.Join(
		stateDir,
		fmt.Sprintf(
			"%s-%x.json",
			strings.TrimLeft(stateRe.ReplaceAllString(logGroupName+"_"+logStreamNamePrefix, "_"), "_"),
			md5.Sum([]byte(
				strings.Join(
					[]string{
						os.Getenv("AWS_PROFILE"),
						os.Getenv("AWS_ACCESS_KEY_ID"),
						os.Getenv("AWS_SECRET_ACCESS_KEY"),
						os.Getenv("AWS_REGION"),
						strings.Join(args, " "),
					},
					" ",
				)),
			),
		),
	)
}

func createAWSConfig(opts *logOpts) *aws.Config {
	conf := aws.NewConfig()
	if opts.MaxRetries > 0 {
		return conf.WithMaxRetries(opts.MaxRetries)
	}
	return conf
}

func createService(opts *logOpts) (*cloudwatchlogs.CloudWatchLogs, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}
	return cloudwatchlogs.New(sess, createAWSConfig(opts)), nil
}

type logState struct {
	NextToken *string
	StartTime *int64
}

func (p *awsCloudwatchLogsPlugin) collect(now time.Time) ([]string, error) {
	s, err := p.loadState()
	if err != nil {
		return nil, err
	}
	if s.StartTime == nil || *s.StartTime <= now.Add(-time.Hour).Unix()*1000 {
		s.StartTime = aws.Int64(now.Add(-1*time.Minute).Unix() * 1000)
	}
	var messages []string
	input := &cloudwatchlogs.FilterLogEventsInput{
		StartTime:     s.StartTime,
		LogGroupName:  aws.String(p.LogGroupName),
		NextToken:     s.NextToken,
		FilterPattern: aws.String(p.Pattern),
	}
	if p.LogStreamNamePrefix != "" {
		input.LogStreamNamePrefix = aws.String(p.LogStreamNamePrefix)
	}
	err = p.Service.FilterLogEventsPages(input, func(output *cloudwatchlogs.FilterLogEventsOutput, lastPage bool) bool {
		for _, event := range output.Events {
			messages = append(messages, *event.Message)
			if s.StartTime == nil || *s.StartTime <= *event.Timestamp {
				s.StartTime = aws.Int64(*event.Timestamp + 1)
			}
		}
		s.NextToken = output.NextToken
		if lastPage {
			s.NextToken = nil
		}
		err = p.saveState(s)
		if err != nil {
			return false
		}
		time.Sleep(150 * time.Millisecond)
		return true
	})
	if err != nil {
		return nil, err
	}
	return messages, nil
}

func (p *awsCloudwatchLogsPlugin) loadState() (*logState, error) {
	var s logState
	f, err := os.Open(p.StateFile)
	defer f.Close()
	if err != nil {
		if os.IsNotExist(err) {
			return &s, nil
		}
		return nil, err
	}
	err = json.NewDecoder(f).Decode(&s)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func (p *awsCloudwatchLogsPlugin) saveState(s *logState) error {
	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(s); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(p.StateFile), 0755); err != nil {
		return err
	}
	return atomic.WriteFile(p.StateFile, &buf)
}

func (p *awsCloudwatchLogsPlugin) check(messages []string) *checkers.Checker {
	status := checkers.OK
	msg := fmt.Sprint(len(messages))
	if len(messages) > p.CriticalOver {
		status = checkers.CRITICAL
		msg += " > " + fmt.Sprint(p.CriticalOver)
	} else if len(messages) > p.WarningOver {
		status = checkers.WARNING
		msg += " > " + fmt.Sprint(p.WarningOver)
	}
	msg += " messages for pattern /" + p.Pattern + "/"
	if status != checkers.OK && p.ReturnContent {
		msg += "\n" + strings.Join(messages, "")
	}
	return checkers.NewChecker(status, msg)
}

func (p *awsCloudwatchLogsPlugin) run(now time.Time) *checkers.Checker {
	messages, err := p.collect(now)
	if err != nil {
		return checkers.Unknown(fmt.Sprint(err))
	}
	return p.check(messages)
}

func run(args []string) *checkers.Checker {
	opts := &logOpts{}
	_, err := flags.ParseArgs(opts, args)
	if err != nil {
		os.Exit(1)
	}
	p, err := newCloudwatchLogsPlugin(opts, args)
	if err != nil {
		return checkers.Unknown(fmt.Sprint(err))
	}
	return p.run(time.Now())
}
<|endoftext|>


<commit_before>package v2ray

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/miekg/dns"
	vcore "v2ray.com/core"
	vnet "v2ray.com/core/common/net"
	vsession "v2ray.com/core/common/session"
	vdns "v2ray.com/core/features/dns"

	"github.com/eycorsican/go-tun2socks/core"
	"github.com/eycorsican/go-tun2socks/proxy"
)

func isIPv4(ip net.IP) bool {
	if ip.To4() != nil {
		return true
	}
	return false
}

func isIPv6(ip net.IP) bool {
	// To16() also valid for ipv4, ensure it's not an ipv4 address
	if ip.To4() != nil {
		return false
	}
	if ip.To16() != nil {
		return true
	}
	return false
}

type connEntry struct {
	conn                net.Conn
	target              vnet.Destination
	cancelFetchingInput context.CancelFunc
	fetchingInputCtx    context.Context
}

type dnsRespEntry struct {
	err  error
	data []byte
	conn core.Connection
}

type handler struct {
	sync.Mutex

	ctx        context.Context
	v          *vcore.Instance
	conns      map[core.Connection]*connEntry
	dispatched map[core.Connection]bool
	dnsRespCh  chan *dnsRespEntry
	dnsClient  vdns.Client
}

func (h *handler) shouldAcceptDNSQuery(data []byte) bool {
	req := new(dns.Msg)
	err := req.Unpack(data)
	if err != nil {
		return false
	}

	// TODO: allow multiple question
	if len(req.Question) != 1 {
		return false
	}

	qtype := req.Question[0].Qtype
	if qtype != dns.TypeA && qtype != dns.TypeAAAA {
		return false
	}

	qclass := req.Question[0].Qclass
	if qclass != dns.ClassINET {
		return false
	}

	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	if _, ok := dns.IsDomainName(domain); !ok {
		return false
	}

	return true
}

func (h *handler) handleDNSQuery(conn core.Connection, data []byte) {
	var err error
	var answer []byte = nil
	defer func() {
		h.dnsRespCh <- &dnsRespEntry{conn: conn, data: answer, err: err}
	}()

	// No error checks here because they are already done in shouldAcceptDNSQuery()
	req := new(dns.Msg)
	req.Unpack(data)
	qtype := req.Question[0].Qtype
	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	log.Printf("dispatch dns request for domain: %v (%v)", domain, qtype)
	ips, err := h.dnsClient.LookupIP(domain)
	if err != nil {
		err = errors.New(fmt.Sprintf("lookup ip failed: %v", err))
		return
	}

	resp := new(dns.Msg)
	resp = resp.SetReply(req)
	resp.RecursionAvailable = true
	for _, ip := range ips {
		if isIPv4(ip) && qtype == dns.TypeA {
			resp.Answer = append(resp.Answer, &dns.A{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 4,
				},
				A: ip,
			})
		} else if isIPv6(ip) && qtype == dns.TypeAAAA {
			resp.Answer = append(resp.Answer, &dns.AAAA{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeAAAA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 16,
				},
				AAAA: ip,
			})
		}
	}
	if len(resp.Answer) == 0 {
		// Has A records but no wanted AAAA records
		if qtype == dns.TypeAAAA && len(ips) != 0 {
			// https://tools.ietf.org/html/rfc4074#section-3
			resp = resp.SetRcode(req, dns.RcodeSuccess)
		} else {
			err = errors.New(fmt.Sprintf("no answer for %v (%v) (%v)", domain, qtype, len(ips)))
			return
		}
	}
	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)
	dnsAnswer, err := resp.PackBuffer(buf)
	if err != nil {
		err = errors.New(fmt.Sprintf("packing dns resp msg failed: %v", err))
	} else {
		answer = append([]byte(nil), dnsAnswer...)
	}
	return
}

func (h *handler) handleDNSResponse() {
	for {
		select {
		case respEntry := <-h.dnsRespCh:
			if respEntry.err == nil {
				_, err := respEntry.conn.Write(respEntry.data)
				if err != nil {
					log.Printf("write dns response to local failed: %v", err)
				}
			} else {
				log.Printf("dispatch dns request failed: %v", respEntry.err)
			}
			h.Close(respEntry.conn)
			respEntry.conn.Close()
		}
	}
}

func (h *handler) fetchInput(conn core.Connection) {
	h.Lock()
	c, ok := h.conns[conn]
	h.Unlock()
	if !ok {
		return
	}

	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)

	for {
		n, err := c.conn.Read(buf)
		select {
		case <-c.fetchingInputCtx.Done():
			// Request was re-dipatched to V2Ray, stop fetching but leave the
			// connection open.
			return
		default:
		}
		select {
		case <-c.fetchingInputCtx.Done():
			// Request was re-dipatched to V2Ray, stop fetching but leave the
			// connection open.
			return
		default:
			if err != nil && n <= 0 {
				h.Close(conn)
				conn.Close()
				return
			}
			_, err = conn.Write(buf[:n])
			if err != nil {
				h.Close(conn)
				conn.Close()
				return
			}
		}
	}
}

func NewHandler(ctx context.Context, instance *vcore.Instance) core.ConnectionHandler {
	h := &handler{
		ctx:        ctx,
		v:          instance,
		conns:      make(map[core.Connection]*connEntry, 16),
		dnsRespCh:  make(chan *dnsRespEntry, 1024),
		dispatched: make(map[core.Connection]bool, 16),
		dnsClient:  instance.GetFeature(vdns.ClientType()).(vdns.Client),
	}
	go h.handleDNSResponse()
	return h
}

func (h *handler) Connect(conn core.Connection, target net.Addr) error {
	dest := vnet.DestinationFromAddr(target)
	sid := vsession.NewID()
	ctx := vsession.ContextWithID(h.ctx, sid)
	c, err := vcore.Dial(ctx, h.v, dest)
	if err != nil {
		return errors.New(fmt.Sprintf("dial V proxy connection failed: %v", err))
	}
	// Note that ctx here is used for canceling fetching input goroutine, not
	// canceling the connection, thus create the cancelable context after Dial().
	ctx, cancel := context.WithCancel(ctx)
	h.Lock()
	h.conns[conn] = &connEntry{
		conn:                c,
		target:              dest,
		cancelFetchingInput: cancel,
		fetchingInputCtx:    ctx,
	}
	h.Unlock()
	go h.fetchInput(conn)
	log.Printf("new proxy connection for target: %s:%s", target.Network(), target.String())
	return nil
}

func (h *handler) DidReceive(conn core.Connection, data []byte) error {
	h.Lock()
	c, ok := h.conns[conn]
	done, ok2 := h.dispatched[conn]
	h.Unlock()
	if ok2 && done {
		// Request already dispatched to V2Ray, ignore.
		return nil
	}
	if ok {
		// If it's a DNS request of type A or AAAA and has only one question,
		// handle it with V2Ray's DNS client, otherwise treat as normal TCP/UDP
		// traffic.
		if c.target.Network == vnet.Network_UDP &&
			c.target.Port.Value() == proxy.COMMON_DNS_PORT &&
			h.shouldAcceptDNSQuery(data) {

			// Parse DNS request and hand to V2Ray, upon V2Ray returns []net.IP,
			// packing them into dns.Msg response message and write back to the client.
			go h.handleDNSQuery(conn, append([]byte(nil), data...))

			// The DNS request has passed to V2Ray for handling, we are safe
			// to cancel the fetching goroutine, but be careful do not close the
			// connection.
			c.cancelFetchingInput()
			c.conn.Close()

			h.Lock()
			// The request is successfully handed to V2Ray, mark as dispatched so
			// subsequent retransmissions should be ignored.
			h.dispatched[conn] = true
			h.Unlock()
		} else {
			_, err := c.conn.Write(data)
			if err != nil {
				h.Close(conn)
				return errors.New(fmt.Sprintf("write remote failed: %v", err))
			}
		}
		return nil
	} else {
		h.Close(conn)
		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), conn.RemoteAddr()))
	}
}

func (h *handler) DidSend(conn core.Connection, len uint16) {
	// unused
}

func (h *handler) DidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) LocalDidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) Close(conn core.Connection) {
	h.Lock()
	defer h.Unlock()

	if c, found := h.conns[conn]; found {
		c.conn.Close()
	}
	delete(h.conns, conn)
}
<commit_msg>remove unused import<commit_after>package v2ray

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"

	"github.com/miekg/dns"
	vcore "v2ray.com/core"
	vnet "v2ray.com/core/common/net"
	vsession "v2ray.com/core/common/session"
	vdns "v2ray.com/core/features/dns"

	"github.com/eycorsican/go-tun2socks/core"
	"github.com/eycorsican/go-tun2socks/proxy"
)

func isIPv4(ip net.IP) bool {
	if ip.To4() != nil {
		return true
	}
	return false
}

func isIPv6(ip net.IP) bool {
	// To16() also valid for ipv4, ensure it's not an ipv4 address
	if ip.To4() != nil {
		return false
	}
	if ip.To16() != nil {
		return true
	}
	return false
}

type connEntry struct {
	conn                net.Conn
	target              vnet.Destination
	cancelFetchingInput context.CancelFunc
	fetchingInputCtx    context.Context
}

type dnsRespEntry struct {
	err  error
	data []byte
	conn core.Connection
}

type handler struct {
	sync.Mutex

	ctx        context.Context
	v          *vcore.Instance
	conns      map[core.Connection]*connEntry
	dispatched map[core.Connection]bool
	dnsRespCh  chan *dnsRespEntry
	dnsClient  vdns.Client
}

func (h *handler) shouldAcceptDNSQuery(data []byte) bool {
	req := new(dns.Msg)
	err := req.Unpack(data)
	if err != nil {
		return false
	}

	// TODO: allow multiple question
	if len(req.Question) != 1 {
		return false
	}

	qtype := req.Question[0].Qtype
	if qtype != dns.TypeA && qtype != dns.TypeAAAA {
		return false
	}

	qclass := req.Question[0].Qclass
	if qclass != dns.ClassINET {
		return false
	}

	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	if _, ok := dns.IsDomainName(domain); !ok {
		return false
	}

	return true
}

func (h *handler) handleDNSQuery(conn core.Connection, data []byte) {
	var err error
	var answer []byte = nil
	defer func() {
		h.dnsRespCh <- &dnsRespEntry{conn: conn, data: answer, err: err}
	}()

	// No error checks here because they are already done in shouldAcceptDNSQuery()
	req := new(dns.Msg)
	req.Unpack(data)
	qtype := req.Question[0].Qtype
	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	log.Printf("dispatch dns request for domain: %v (%v)", domain, qtype)
	ips, err := h.dnsClient.LookupIP(domain)
	if err != nil {
		err = errors.New(fmt.Sprintf("lookup ip failed: %v", err))
		return
	}

	resp := new(dns.Msg)
	resp = resp.SetReply(req)
	resp.RecursionAvailable = true
	for _, ip := range ips {
		if isIPv4(ip) && qtype == dns.TypeA {
			resp.Answer = append(resp.Answer, &dns.A{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 4,
				},
				A: ip,
			})
		} else if isIPv6(ip) && qtype == dns.TypeAAAA {
			resp.Answer = append(resp.Answer, &dns.AAAA{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeAAAA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 16,
				},
				AAAA: ip,
			})
		}
	}
	if len(resp.Answer) == 0 {
		// Has A records but no wanted AAAA records
		if qtype == dns.TypeAAAA && len(ips) != 0 {
			// https://tools.ietf.org/html/rfc4074#section-3
			resp = resp.SetRcode(req, dns.RcodeSuccess)
		} else {
			err = errors.New(fmt.Sprintf("no answer for %v (%v) (%v)", domain, qtype, len(ips)))
			return
		}
	}
	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)
	dnsAnswer, err := resp.PackBuffer(buf)
	if err != nil {
		err = errors.New(fmt.Sprintf("packing dns resp msg failed: %v", err))
	} else {
		answer = append([]byte(nil), dnsAnswer...)
	}
	return
}

func (h *handler) handleDNSResponse() {
	for {
		select {
		case respEntry := <-h.dnsRespCh:
			if respEntry.err == nil {
				_, err := respEntry.conn.Write(respEntry.data)
				if err != nil {
					log.Printf("write dns response to local failed: %v", err)
				}
			} else {
				log.Printf("dispatch dns request failed: %v", respEntry.err)
			}
			h.Close(respEntry.conn)
			respEntry.conn.Close()
		}
	}
}

func (h *handler) fetchInput(conn core.Connection) {
	h.Lock()
	c, ok := h.conns[conn]
	h.Unlock()
	if !ok {
		return
	}

	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)

	for {
		n, err := c.conn.Read(buf)
		select {
		case <-c.fetchingInputCtx.Done():
			// Request was re-dipatched to V2Ray, stop fetching but leave the
			// connection open.
			return
		default:
		}
		select {
		case <-c.fetchingInputCtx.Done():
			// Request was re-dipatched to V2Ray, stop fetching but leave the
			// connection open.
			return
		default:
			if err != nil && n <= 0 {
				h.Close(conn)
				conn.Close()
				return
			}
			_, err = conn.Write(buf[:n])
			if err != nil {
				h.Close(conn)
				conn.Close()
				return
			}
		}
	}
}

func NewHandler(ctx context.Context, instance *vcore.Instance) core.ConnectionHandler {
	h := &handler{
		ctx:        ctx,
		v:          instance,
		conns:      make(map[core.Connection]*connEntry, 16),
		dnsRespCh:  make(chan *dnsRespEntry, 1024),
		dispatched: make(map[core.Connection]bool, 16),
		dnsClient:  instance.GetFeature(vdns.ClientType()).(vdns.Client),
	}
	go h.handleDNSResponse()
	return h
}

func (h *handler) Connect(conn core.Connection, target net.Addr) error {
	dest := vnet.DestinationFromAddr(target)
	sid := vsession.NewID()
	ctx := vsession.ContextWithID(h.ctx, sid)
	c, err := vcore.Dial(ctx, h.v, dest)
	if err != nil {
		return errors.New(fmt.Sprintf("dial V proxy connection failed: %v", err))
	}
	// Note that ctx here is used for canceling fetching input goroutine, not
	// canceling the connection, thus create the cancelable context after Dial().
	ctx, cancel := context.WithCancel(ctx)
	h.Lock()
	h.conns[conn] = &connEntry{
		conn:                c,
		target:              dest,
		cancelFetchingInput: cancel,
		fetchingInputCtx:    ctx,
	}
	h.Unlock()
	go h.fetchInput(conn)
	log.Printf("new proxy connection for target: %s:%s", target.Network(), target.String())
	return nil
}

func (h *handler) DidReceive(conn core.Connection, data []byte) error {
	h.Lock()
	c, ok := h.conns[conn]
	done, ok2 := h.dispatched[conn]
	h.Unlock()
	if ok2 && done {
		// Request already dispatched to V2Ray, ignore.
		return nil
	}
	if ok {
		// If it's a DNS request of type A or AAAA and has only one question,
		// handle it with V2Ray's DNS client, otherwise treat as normal TCP/UDP
		// traffic.
		if c.target.Network == vnet.Network_UDP &&
			c.target.Port.Value() == proxy.COMMON_DNS_PORT &&
			h.shouldAcceptDNSQuery(data) {

			// Parse DNS request and hand to V2Ray, upon V2Ray returns []net.IP,
			// packing them into dns.Msg response message and write back to the client.
			go h.handleDNSQuery(conn, append([]byte(nil), data...))

			// The DNS request has passed to V2Ray for handling, we are safe
			// to cancel the fetching goroutine, but be careful do not close the
			// connection.
			c.cancelFetchingInput()
			c.conn.Close()

			h.Lock()
			// The request is successfully handed to V2Ray, mark as dispatched so
			// subsequent retransmissions should be ignored.
			h.dispatched[conn] = true
			h.Unlock()
		} else {
			_, err := c.conn.Write(data)
			if err != nil {
				h.Close(conn)
				return errors.New(fmt.Sprintf("write remote failed: %v", err))
			}
		}
		return nil
	} else {
		h.Close(conn)
		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), conn.RemoteAddr()))
	}
}

func (h *handler) DidSend(conn core.Connection, len uint16) {
	// unused
}

func (h *handler) DidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) LocalDidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) Close(conn core.Connection) {
	h.Lock()
	defer h.Unlock()

	if c, found := h.conns[conn]; found {
		c.conn.Close()
	}
	delete(h.conns, conn)
}
<|endoftext|>


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

// Package suite allows to read tests and collections of tests (suites) from
// disk and execute them in a controlled way or run throughput load test from
// these test/suites.
//
package suite

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"sort"
	"strings"
	"time"

	"github.com/vdobler/ht/cookiejar"
	"github.com/vdobler/ht/ht"
	"github.com/vdobler/ht/mock"
	"github.com/vdobler/ht/scope"
)

// A Suite is a collection of Tests which can be executed sequentily with the
// result captured.
type Suite struct {
	Name        string // Name of the Suite.
	Description string // Description of what's going on here.
	KeepCookies bool   // KeepCookies in a cookie jar common to all Tests.

	Status   ht.Status     // Status is the overall status of the whole suite.
	Error    error         // Error encountered during execution of the suite.
	Started  time.Time     // Start of the execution.
	Duration time.Duration // Duration of the execution.

	Tests []*ht.Test // The Tests to execute

	Variables      scope.Variables // The initial variable assignment
	FinalVariables scope.Variables // The final set of variables.
	Jar            *cookiejar.Jar  // The cookie jar used

	Verbosity int
	Log       interface {
		Printf(format string, a ...interface{})
	}

	globals          scope.Variables
	tests            []*RawTest
	noneTeardownTest int
}

func shouldRun(t int, rs *RawSuite, s *Suite) bool {
	if !rs.tests[t].IsEnabled() {
		return false
	}

	// Stop execution on errors during setup
	for i := 0; i < len(rs.Setup) && i < len(s.Tests); i++ {
		if s.Tests[i].Status > ht.Pass {
			return false
		}
	}
	return true
}

// NewFromRaw sets up a new Suite from rs, read to be Iterated.
func NewFromRaw(rs *RawSuite, global map[string]string, jar *cookiejar.Jar, logger *log.Logger) *Suite {
	// Create cookie jar if needed.
	if rs.KeepCookies {
		if jar == nil {
			// Make own, private-use jar.
			jar, _ = cookiejar.New(nil)
		}
	} else {
		jar = nil
	}

	if logger == nil {
		logger = log.New(ioutil.Discard, "", 0)
	}

	suite := &Suite{
		KeepCookies: rs.KeepCookies,

		Status: ht.NotRun,
		Error:  nil,

		Tests: make([]*ht.Test, 0, len(rs.tests)),

		Variables:        make(map[string]string),
		FinalVariables:   make(map[string]string),
		Jar:              jar,
		Log:              logger,
		Verbosity:        rs.Verbosity,
		tests:            rs.tests,
		noneTeardownTest: len(rs.Setup) + len(rs.Main),
	}

	suite.globals = scope.New(global, rs.Variables, true)
	suite.globals["SUITE_DIR"] = rs.File.Dirname()
	suite.globals["SUITE_NAME"] = rs.File.Basename()
	replacer := suite.globals.Replacer()

	suite.Name = replacer.Replace(rs.Name)
	suite.Description = replacer.Replace(rs.Description)

	for n, v := range suite.globals {
		suite.Variables[n] = v
	}

	return suite
}

// A Executor is responsible for executing the given test during the
// Iterate'ion of a Suite. It should return nil if execution should continue
// and ErrAbortExecution to stop further iteration.
type Executor func(test *ht.Test) error

var (
	// ErrAbortExecution indicates that suite iteration should stop.
	ErrAbortExecution = errors.New("Abort Execution")
)

var mockDelay = 50 * time.Millisecond

// Iterate the suite through the given executor.
func (suite *Suite) Iterate(executor Executor) {
	now := time.Now()
	now = now.Add(-time.Duration(now.Nanosecond()))
	suite.Started = now

	overall := ht.NotRun
	errors := ht.ErrorList{}

	for _, rt := range suite.tests {
		// suite.Log.Printf("Executing Test %q\n", rt.File.Name)
		callScope := scope.New(suite.globals, rt.contextVars, true)
		testScope := scope.New(callScope, rt.Variables, false)
		testScope["TEST_DIR"] = rt.File.Dirname()
		testScope["TEST_NAME"] = rt.File.Basename()
		test, err := rt.ToTest(testScope)
		if err != nil {
			test.Status = ht.Bogus
			test.Error = err
		}
		test.Jar = suite.Jar
		test.Log = suite.Log

		// Mocks requested for this test: We expect each mock to be
		// called exactly once (and this call should pass).
		var stopMocks, monitoringDone chan bool
		var monitor chan *ht.Test
		var mockResult []*ht.Test
		var expect []string
		if len(rt.mocks) != 0 {
			monitor = make(chan *ht.Test)
			mocks := []*mock.Mock{}
			for i, m := range rt.mocks {
				mockScope := scope.New(testScope, rt.Variables, false)
				mockScope["MOCK_DIR"] = m.Dirname()
				mockScope["MOCK_NAME"] = m.Basename()
				if mk, err := m.ToMock(mockScope, true); err != nil {
					test.Status = ht.Bogus
					test.Error = fmt.Errorf(
						"Mock %d %q is malformed: %s",
						i+1, m.Name, err)
				} else {
					mk.Monitor = monitor
					mocks = append(mocks, mk)
					expect = append(expect, mk.Name)
				}
			}
			if len(mocks) > 0 {
				// TODO handle 404s
				stopMocks, err = mock.Serve(mocks, nil, suite.Log)
				// TODO: handle error
				monitoringDone = make(chan bool)
				go func() {
					for report := range monitor {
						logMock(suite, report)
						mockResult = append(mockResult, report)
					}
					close(monitoringDone)
				}()
				time.Sleep(mockDelay) // I'm clueless...
			}
		}

		exstat := executor(test)
		if stopMocks != nil {
			stopMocks <- true
			<-stopMocks
		}
		if monitor != nil {
			close(monitor)
		}
		if monitoringDone != nil {
			<-monitoringDone
			// Analyse what we got.
			actual := []string{}
			for _, mt := range mockResult {
				actual = append(actual, mt.Name)
				// Propagete state of mock invocations to main test.
				if mt.Status > test.Status {
					test.Status = mt.Status
					test.Error = mt.Error
				}
				// Augment check name with information about mock.
				for i := range mt.CheckResults {
					name := mt.CheckResults[i].Name
					name = fmt.Sprintf("Request to mock %q: %s", mt.Name, name)
					mt.CheckResults[i].Name = name
				}
				// Append check results of mock to main test checks.
				test.CheckResults = append(test.CheckResults, mt.CheckResults...)
			}

			// Did we get exactly what we expected?
			sort.Strings(expect)
			sort.Strings(actual)
			// Sorry for that.
			want, got := strings.Join(expect, " ※ "), strings.Join(actual, " ※ ")
			if got != want && test.Status == ht.Pass {
				test.Status = ht.Fail
				test.Error = fmt.Errorf("Expected mocks [%s] but actual mock invocations were [%s]",
					want, got)
			}
		}
		if test.Status == ht.Pass {
			suite.updateVariables(test)
		}

		suite.Tests = append(suite.Tests, test)
		if test.Status > overall {
			overall = test.Status
		}
		if err := test.Error; err != nil {
			errors = append(errors, err)
		}

		if exstat == ErrAbortExecution {
			break
		}
	}
	suite.Duration = time.Since(suite.Started)
	clip := suite.Duration.Nanoseconds() % 1000000
	suite.Duration -= time.Duration(clip)
	suite.Status = overall
	if len(errors) == 0 {
		suite.Error = nil
	} else {
		suite.Error = errors
	}

	for n, v := range suite.globals {
		suite.FinalVariables[n] = v
	}
}

func logMock(suite *Suite, report *ht.Test) {
	if suite.Verbosity <= 0 {
		return
	}
	full := ""
	if suite.Verbosity >= 3 {
		full = "\n"
		full += fmt.Sprintf("  Request\n    Header\n")
		for k, v := range report.Request.Request.Header {
			full += fmt.Sprintf("      %s: %s\n", k, v)
		}
		full += fmt.Sprintf("    Body\n")
		full += fmt.Sprintf("      %s\n", report.Request.SentBody)
		full += fmt.Sprintf("  Response\n    Header\n")
		for k, v := range report.Response.Response.Header {
			full += fmt.Sprintf("      %s: %s\n", k, v)
		}
		full += fmt.Sprintf("    Body\n")
		full += fmt.Sprintf("      %s\n", report.Response.BodyStr)
		full += fmt.Sprintf("========================================================\n")

	}
	suite.Log.Printf("Mock invoked %q: %s %s%s", report.Name,
		report.Request.Method, report.Request.URL, full)

}

func (suite *Suite) updateVariables(test *ht.Test) {
	if test.Status != ht.Pass {
		return
	}

	for varname, value := range test.Extract() {
		if suite.Verbosity >= 2 {
			if old, ok := suite.globals[varname]; ok {
				if value != old {
					suite.Log.Printf("Updating variable %q to %q\n",
						varname, value)
				} else {
					suite.Log.Printf("Keeping  variable %q as %q\n",
						varname, value)
				}
			} else {
				suite.Log.Printf("Setting  variable %q to %q\n",
					varname, value)
			}
		}

		suite.globals[varname] = value
	}
}

func (suite *Suite) updateStatus(test *ht.Test) {
	if test.Status <= suite.Status {
		return
	}

	suite.Status = test.Status
	if test.Error != nil {
		suite.Error = test.Error
	}
}

// Stats counts the test results of s.
func (suite *Suite) Stats() (notRun int, skipped int, passed int, failed int, errored int, bogus int) {
	for _, tr := range suite.Tests {
		switch tr.Status {
		case ht.NotRun:
			notRun++
		case ht.Skipped:
			skipped++
		case ht.Pass:
			passed++
		case ht.Fail:
			failed++
		case ht.Error:
			errored++
		case ht.Bogus:
			bogus++
		default:
			panic(fmt.Sprintf("No such Status %d in suite %q test %q",
				tr.Status, suite.Name, tr.Name))
		}
	}
	return
}
<commit_msg>suite: refacator suite execution<commit_after>// Copyright 2016 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package suite allows to read tests and collections of tests (suites) from
// disk and execute them in a controlled way or run throughput load test from
// these test/suites.
//
package suite

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"sort"
	"strings"
	"time"

	"github.com/vdobler/ht/cookiejar"
	"github.com/vdobler/ht/ht"
	"github.com/vdobler/ht/mock"
	"github.com/vdobler/ht/scope"
)

// A Suite is a collection of Tests which can be executed sequentily with the
// result captured.
type Suite struct {
	Name        string // Name of the Suite.
	Description string // Description of what's going on here.
	KeepCookies bool   // KeepCookies in a cookie jar common to all Tests.

	Status   ht.Status     // Status is the overall status of the whole suite.
	Error    error         // Error encountered during execution of the suite.
	Started  time.Time     // Start of the execution.
	Duration time.Duration // Duration of the execution.

	Tests []*ht.Test // The Tests to execute

	Variables      scope.Variables // The initial variable assignment
	FinalVariables scope.Variables // The final set of variables.
	Jar            *cookiejar.Jar  // The cookie jar used

	Verbosity int
	Log       interface {
		Printf(format string, a ...interface{})
	}

	globals          scope.Variables
	tests            []*RawTest
	noneTeardownTest int
}

func shouldRun(t int, rs *RawSuite, s *Suite) bool {
	if !rs.tests[t].IsEnabled() {
		return false
	}

	// Stop execution on errors during setup
	for i := 0; i < len(rs.Setup) && i < len(s.Tests); i++ {
		if s.Tests[i].Status > ht.Pass {
			return false
		}
	}
	return true
}

// NewFromRaw sets up a new Suite from rs, read to be Iterated.
func NewFromRaw(rs *RawSuite, global map[string]string, jar *cookiejar.Jar, logger *log.Logger) *Suite {
	// Create cookie jar if needed.
	if rs.KeepCookies {
		if jar == nil {
			// Make own, private-use jar.
			jar, _ = cookiejar.New(nil)
		}
	} else {
		jar = nil
	}

	if logger == nil {
		logger = log.New(ioutil.Discard, "", 0)
	}

	suite := &Suite{
		KeepCookies: rs.KeepCookies,

		Status: ht.NotRun,
		Error:  nil,

		Tests: make([]*ht.Test, 0, len(rs.tests)),

		Variables:        make(map[string]string),
		FinalVariables:   make(map[string]string),
		Jar:              jar,
		Log:              logger,
		Verbosity:        rs.Verbosity,
		tests:            rs.tests,
		noneTeardownTest: len(rs.Setup) + len(rs.Main),
	}

	suite.globals = scope.New(global, rs.Variables, true)
	suite.globals["SUITE_DIR"] = rs.File.Dirname()
	suite.globals["SUITE_NAME"] = rs.File.Basename()
	replacer := suite.globals.Replacer()

	suite.Name = replacer.Replace(rs.Name)
	suite.Description = replacer.Replace(rs.Description)

	for n, v := range suite.globals {
		suite.Variables[n] = v
	}

	return suite
}

// A Executor is responsible for executing the given test during the
// Iterate'ion of a Suite. It should return nil if execution should continue
// and ErrAbortExecution to stop further iteration.
type Executor func(test *ht.Test) error

var (
	// ErrAbortExecution indicates that suite iteration should stop.
	ErrAbortExecution = errors.New("Abort Execution")
)

var mockDelay = 50 * time.Millisecond

// Iterate the suite through the given executor.
func (suite *Suite) Iterate(executor Executor) {
	now := time.Now()
	now = now.Add(-time.Duration(now.Nanosecond()))
	suite.Started = now

	overall := ht.NotRun
	errors := ht.ErrorList{}

	for _, rt := range suite.tests {
		// suite.Log.Printf("Executing Test %q\n", rt.File.Name)
		callScope := scope.New(suite.globals, rt.contextVars, true)
		testScope := scope.New(callScope, rt.Variables, false)
		testScope["TEST_DIR"] = rt.File.Dirname()
		testScope["TEST_NAME"] = rt.File.Basename()
		test, err := rt.ToTest(testScope)
		if err != nil {
			test.Status = ht.Bogus
			test.Error = err
		}
		test.Jar = suite.Jar
		test.Log = suite.Log

		// Mocks requested for this test: We expect each mock to be
		// called exactly once (and this call should pass).
		var mockResult []*ht.Test
		stopMocks, monitor, monitoringDone, expect, err := startMocks(suite, test, rt, &mockResult, testScope)
		if err != nil {
			test.Status = ht.Bogus
			test.Error = err
		}

		// Execute the test (if not bogus).
		exstat := executor(test)

		if stopMocks != nil {
			// We got running mocks: Stop mock handling and stop monitoring
			stopMocks <- true
			<-stopMocks
			close(monitor)
			<-monitoringDone

			// Analyse what we got and updates test.
			analyseMocks(test, mockResult, expect)
		}
		if test.Status == ht.Pass {
			suite.updateVariables(test)
		}

		suite.Tests = append(suite.Tests, test)
		if test.Status > overall {
			overall = test.Status
		}
		if err := test.Error; err != nil {
			errors = append(errors, err)
		}

		if exstat == ErrAbortExecution {
			break
		}
	}
	suite.Duration = time.Since(suite.Started)
	clip := suite.Duration.Nanoseconds() % 1000000
	suite.Duration -= time.Duration(clip)
	suite.Status = overall
	if len(errors) == 0 {
		suite.Error = nil
	} else {
		suite.Error = errors
	}

	for n, v := range suite.globals {
		suite.FinalVariables[n] = v
	}
}

func startMocks(suite *Suite, test *ht.Test, rt *RawTest, mockResult *[]*ht.Test, testScope scope.Variables) (stopMocks chan bool, monitor chan *ht.Test, monitoringDone chan bool, expect []string, err error) {
	if len(rt.mocks) == 0 {
		return nil, nil, nil, nil, nil
	}

	monitor = make(chan *ht.Test)
	mocks := []*mock.Mock{}
	for i, m := range rt.mocks {
		mockScope := scope.New(testScope, rt.Variables, false)
		mockScope["MOCK_DIR"] = m.Dirname()
		mockScope["MOCK_NAME"] = m.Basename()
		mk, err := m.ToMock(mockScope, true)
		if err != nil {
			return nil, nil, nil, nil,
				fmt.Errorf("mock %d %q is malformed: %s",
					i+1, m.Name, err)

		}
		mk.Monitor = monitor
		mocks = append(mocks, mk)
		expect = append(expect, mk.Name)
	}

	// TODO: handle 404
	stopMocks, err = mock.Serve(mocks, nil, suite.Log)
	if err != nil {
		return nil, nil, nil, nil, err
	}

	monitoringDone = make(chan bool)
	go func() {
		for report := range monitor {
			logMock(suite, report)
			*mockResult = append(*mockResult, report)
		}
		close(monitoringDone)
	}()
	time.Sleep(mockDelay) // I'm clueless...

	return stopMocks, monitor, monitoringDone, expect, nil
}

func analyseMocks(test *ht.Test, mockResult []*ht.Test, expect []string) {
	actual := []string{}
	for _, mt := range mockResult {
		actual = append(actual, mt.Name)
		// Propagete state of mock invocations to main test.
		if mt.Status > test.Status {
			test.Status = mt.Status
			test.Error = mt.Error
		}
		// Augment check name with information about mock.
		for i := range mt.CheckResults {
			name := mt.CheckResults[i].Name
			name = fmt.Sprintf("Request to mock %q: %s", mt.Name, name)
			mt.CheckResults[i].Name = name
		}
		// Append check results of mock to main test checks.
		test.CheckResults = append(test.CheckResults, mt.CheckResults...)
	}

	// Did we get exactly what we expected?
	sort.Strings(expect)
	sort.Strings(actual)
	// Sorry for that.
	want, got := strings.Join(expect, " ※ "), strings.Join(actual, " ※ ")
	if got != want && test.Status == ht.Pass {
		test.Status = ht.Fail
		test.Error = fmt.Errorf("Expected mocks [%s] but actual mock invocations were [%s]",
			want, got)
	}
}

func logMock(suite *Suite, report *ht.Test) {
	if suite.Verbosity <= 0 {
		return
	}
	full := ""
	if suite.Verbosity >= 3 {
		full = "\n"
		full += fmt.Sprintf("  Request\n    Header\n")
		for k, v := range report.Request.Request.Header {
			full += fmt.Sprintf("      %s: %s\n", k, v)
		}
		full += fmt.Sprintf("    Body\n")
		full += fmt.Sprintf("      %s\n", report.Request.SentBody)
		full += fmt.Sprintf("  Response\n    Header\n")
		for k, v := range report.Response.Response.Header {
			full += fmt.Sprintf("      %s: %s\n", k, v)
		}
		full += fmt.Sprintf("    Body\n")
		full += fmt.Sprintf("      %s\n", report.Response.BodyStr)
		full += fmt.Sprintf("========================================================\n")

	}
	suite.Log.Printf("Mock invoked %q: %s %s%s", report.Name,
		report.Request.Method, report.Request.URL, full)

}

func (suite *Suite) updateVariables(test *ht.Test) {
	if test.Status != ht.Pass {
		return
	}

	for varname, value := range test.Extract() {
		if suite.Verbosity >= 2 {
			if old, ok := suite.globals[varname]; ok {
				if value != old {
					suite.Log.Printf("Updating variable %q to %q\n",
						varname, value)
				} else {
					suite.Log.Printf("Keeping  variable %q as %q\n",
						varname, value)
				}
			} else {
				suite.Log.Printf("Setting  variable %q to %q\n",
					varname, value)
			}
		}

		suite.globals[varname] = value
	}
}

func (suite *Suite) updateStatus(test *ht.Test) {
	if test.Status <= suite.Status {
		return
	}

	suite.Status = test.Status
	if test.Error != nil {
		suite.Error = test.Error
	}
}

// Stats counts the test results of s.
func (suite *Suite) Stats() (notRun int, skipped int, passed int, failed int, errored int, bogus int) {
	for _, tr := range suite.Tests {
		switch tr.Status {
		case ht.NotRun:
			notRun++
		case ht.Skipped:
			skipped++
		case ht.Pass:
			passed++
		case ht.Fail:
			failed++
		case ht.Error:
			errored++
		case ht.Bogus:
			bogus++
		default:
			panic(fmt.Sprintf("No such Status %d in suite %q test %q",
				tr.Status, suite.Name, tr.Name))
		}
	}
	return
}
<|endoftext|>


<commit_before>package images

import (
	"Fantasia/system"
	"errors"
	"log"

	"github.com/bwmarrin/discordgo"
)

//genmodules:config

// errors
var (
	ErrNoImagesFound = errors.New("No images found")
)

// MessageCacheLimit sets the cache limit of the images chache
const MessageCacheLimit = 10

// Config ...
type Config struct {
}

// NewConfig returns a pointer to a new config
func NewConfig() *Config {
	return &Config{}
}

// Module ...
type Module struct {
	Sys      *system.System
	Config   *Config
	ImgCache *MessageCache
}

// Build builds the module
func (m *Module) Build(sys *system.System) {
	m.Sys = sys

	// Create Cache
	m.ImgCache = NewMessageCache(MessageCacheLimit)

	// Create image commands
	m.CreateCommands()

	// Add messages with images to the state
	go m.TrackImages()
}

// TrackImages tracks messages that are images and inserts them into the cache
func (m *Module) TrackImages() {
	m.Sys.Dream.DG.AddHandler(func(_ *discordgo.Session, msg *discordgo.MessageCreate) {
		if HasImage(msg.Message) {
			err := m.ImgCache.Add(msg.ChannelID, msg.Message)
			if err != nil {
				log.Println("images: error adding message to cache: ", err)
			}
		}
	})
}
<commit_msg>I don't need a goroutine here<commit_after>package images

import (
	"Fantasia/system"
	"errors"
	"log"

	"github.com/bwmarrin/discordgo"
)

//genmodules:config

// errors
var (
	ErrNoImagesFound = errors.New("No images found")
)

// MessageCacheLimit sets the cache limit of the images chache
const MessageCacheLimit = 10

// Config ...
type Config struct {
}

// NewConfig returns a pointer to a new config
func NewConfig() *Config {
	return &Config{}
}

// Module ...
type Module struct {
	Sys      *system.System
	Config   *Config
	ImgCache *MessageCache
}

// Build builds the module
func (m *Module) Build(sys *system.System) {
	m.Sys = sys

	// Create Cache
	m.ImgCache = NewMessageCache(MessageCacheLimit)

	// Create image commands
	m.CreateCommands()

	// Add messages with images to the state
	m.TrackImages()
}

// TrackImages tracks messages that are images and inserts them into the cache
func (m *Module) TrackImages() {
	m.Sys.Dream.DG.AddHandler(func(_ *discordgo.Session, msg *discordgo.MessageCreate) {
		if HasImage(msg.Message) {
			err := m.ImgCache.Add(msg.ChannelID, msg.Message)
			if err != nil {
				log.Println("images: error adding message to cache: ", err)
			}
		}
	})
}
<|endoftext|>