<commit_before>package aws

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"sort"
	"strings"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/go-homedir"

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

func resourceAwsS3BucketObject() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsS3BucketObjectPut,
		Read:   resourceAwsS3BucketObjectRead,
		Update: resourceAwsS3BucketObjectPut,
		Delete: resourceAwsS3BucketObjectDelete,

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

			"acl": {
				Type:         schema.TypeString,
				Default:      "private",
				Optional:     true,
				ValidateFunc: validateS3BucketObjectAclType,
			},

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

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

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

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

			"content_type": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"key": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"source": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"content"},
			},

			"content": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"source"},
			},

			"storage_class": {
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				ValidateFunc: validateS3BucketObjectStorageClassType,
			},

			"server_side_encryption": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateS3BucketObjectServerSideEncryption,
				Computed:     true,
			},

			"kms_key_id": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateArn,
			},

			"etag": {
				Type: schema.TypeString,
				// This will conflict with SSE-C and SSE-KMS encryption and multi-part upload
				// if/when it's actually implemented. The Etag then won't match raw-file MD5.
				// See http://docs.aws.amazon.com/AmazonS3/latest/API/RESTCommonResponseHeaders.html
				Optional:      true,
				Computed:      true,
				ConflictsWith: []string{"kms_key_id", "server_side_encryption"},
			},

			"version_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),

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

func resourceAwsS3BucketObjectPut(d *schema.ResourceData, meta interface{}) error {
	s3conn := meta.(*AWSClient).s3conn

	restricted := meta.(*AWSClient).IsGovCloud() || meta.(*AWSClient).IsChinaCloud()

	var body io.ReadSeeker

	if v, ok := d.GetOk("source"); ok {
		source := v.(string)
		path, err := homedir.Expand(source)
		if err != nil {
			return fmt.Errorf("Error expanding homedir in source (%s): %s", source, err)
		}
		file, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Error opening S3 bucket object source (%s): %s", source, err)
		}

		body = file
	} else if v, ok := d.GetOk("content"); ok {
		content := v.(string)
		body = bytes.NewReader([]byte(content))
	} else {
		return fmt.Errorf("Must specify \"source\" or \"content\" field")
	}

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

	putInput := &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
		ACL:    aws.String(d.Get("acl").(string)),
		Body:   body,
	}

	if v, ok := d.GetOk("storage_class"); ok {
		putInput.StorageClass = aws.String(v.(string))
	}

	if v, ok := d.GetOk("cache_control"); ok {
		putInput.CacheControl = aws.String(v.(string))
	}

	if v, ok := d.GetOk("content_type"); ok {
		putInput.ContentType = aws.String(v.(string))
	}

	if v, ok := d.GetOk("content_encoding"); ok {
		putInput.ContentEncoding = aws.String(v.(string))
	}

	if v, ok := d.GetOk("content_language"); ok {
		putInput.ContentLanguage = aws.String(v.(string))
	}

	if v, ok := d.GetOk("content_disposition"); ok {
		putInput.ContentDisposition = aws.String(v.(string))
	}

	if v, ok := d.GetOk("server_side_encryption"); ok {
		putInput.ServerSideEncryption = aws.String(v.(string))
	}

	if v, ok := d.GetOk("kms_key_id"); ok {
		putInput.SSEKMSKeyId = aws.String(v.(string))
		putInput.ServerSideEncryption = aws.String(s3.ServerSideEncryptionAwsKms)
	}

	if v, ok := d.GetOk("tags"); ok {
		if restricted {
			return fmt.Errorf("This region does not allow for tags on S3 objects")
		}

		// The tag-set must be encoded as URL Query parameters.
		values := url.Values{}
		for k, v := range v.(map[string]interface{}) {
			values.Add(k, v.(string))
		}
		putInput.Tagging = aws.String(values.Encode())
	}

	if v, ok := d.GetOk("website_redirect"); ok {
		putInput.WebsiteRedirectLocation = aws.String(v.(string))
	}

	resp, err := s3conn.PutObject(putInput)
	if err != nil {
		return fmt.Errorf("Error putting object in S3 bucket (%s): %s", bucket, err)
	}

	// See https://forums.aws.amazon.com/thread.jspa?threadID=44003
	d.Set("etag", strings.Trim(*resp.ETag, `"`))

	d.Set("version_id", resp.VersionId)
	d.SetId(key)
	return resourceAwsS3BucketObjectRead(d, meta)
}

func resourceAwsS3BucketObjectRead(d *schema.ResourceData, meta interface{}) error {
	s3conn := meta.(*AWSClient).s3conn

	restricted := meta.(*AWSClient).IsGovCloud() || meta.(*AWSClient).IsChinaCloud()

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

	resp, err := s3conn.HeadObject(
		&s3.HeadObjectInput{
			Bucket: aws.String(bucket),
			Key:    aws.String(key),
		})

	if err != nil {
		// If S3 returns a 404 Request Failure, mark the object as destroyed
		if awsErr, ok := err.(awserr.RequestFailure); ok && awsErr.StatusCode() == 404 {
			d.SetId("")
			log.Printf("[WARN] Error Reading Object (%s), object not found (HTTP status 404)", key)
			return nil
		}
		return err
	}
	log.Printf("[DEBUG] Reading S3 Bucket Object meta: %s", resp)

	d.Set("cache_control", resp.CacheControl)
	d.Set("content_disposition", resp.ContentDisposition)
	d.Set("content_encoding", resp.ContentEncoding)
	d.Set("content_language", resp.ContentLanguage)
	d.Set("content_type", resp.ContentType)
	d.Set("version_id", resp.VersionId)
	d.Set("server_side_encryption", resp.ServerSideEncryption)
	d.Set("website_redirect", resp.WebsiteRedirectLocation)

	// Only set non-default KMS key ID (one that doesn't match default)
	if resp.SSEKMSKeyId != nil {
		// retrieve S3 KMS Default Master Key
		kmsconn := meta.(*AWSClient).kmsconn
		kmsresp, err := kmsconn.DescribeKey(&kms.DescribeKeyInput{
			KeyId: aws.String("alias/aws/s3"),
		})
		if err != nil {
			return fmt.Errorf("Failed to describe default S3 KMS key (alias/aws/s3): %s", err)
		}

		if *resp.SSEKMSKeyId != *kmsresp.KeyMetadata.Arn {
			log.Printf("[DEBUG] S3 object is encrypted using a non-default KMS Key ID: %s", *resp.SSEKMSKeyId)
			d.Set("kms_key_id", resp.SSEKMSKeyId)
		}
	}
	d.Set("etag", strings.Trim(*resp.ETag, `"`))

	// The "STANDARD" (which is also the default) storage
	// class when set would not be included in the results.
	d.Set("storage_class", s3.StorageClassStandard)
	if resp.StorageClass != nil {
		d.Set("storage_class", resp.StorageClass)
	}

	if !restricted {
		tagResp, err := s3conn.GetObjectTagging(
			&s3.GetObjectTaggingInput{
				Bucket: aws.String(bucket),
				Key:    aws.String(key),
			})
		if err != nil {
			return fmt.Errorf("Failed to get object tags (bucket: %s, key: %s): %s", bucket, key, err)
		}
		d.Set("tags", tagsToMapS3(tagResp.TagSet))
	}

	return nil
}

func resourceAwsS3BucketObjectDelete(d *schema.ResourceData, meta interface{}) error {
	s3conn := meta.(*AWSClient).s3conn

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

	if _, ok := d.GetOk("version_id"); ok {
		// Bucket is versioned, we need to delete all versions
		vInput := s3.ListObjectVersionsInput{
			Bucket: aws.String(bucket),
			Prefix: aws.String(key),
		}
		out, err := s3conn.ListObjectVersions(&vInput)
		if err != nil {
			return fmt.Errorf("Failed listing S3 object versions: %s", err)
		}

		for _, v := range out.Versions {
			input := s3.DeleteObjectInput{
				Bucket:    aws.String(bucket),
				Key:       aws.String(key),
				VersionId: v.VersionId,
			}
			_, err := s3conn.DeleteObject(&input)
			if err != nil {
				return fmt.Errorf("Error deleting S3 object version of %s:\n %s:\n %s",
					key, v, err)
			}
		}
	} else {
		// Just delete the object
		input := s3.DeleteObjectInput{
			Bucket: aws.String(bucket),
			Key:    aws.String(key),
		}
		_, err := s3conn.DeleteObject(&input)
		if err != nil {
			return fmt.Errorf("Error deleting S3 bucket object: %s  Bucket: %q Object: %q", err, bucket, key)
		}
	}

	return nil
}

func validateS3BucketObjectAclType(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)

	cannedAcls := map[string]bool{
		s3.ObjectCannedACLPrivate:                true,
		s3.ObjectCannedACLPublicRead:             true,
		s3.ObjectCannedACLPublicReadWrite:        true,
		s3.ObjectCannedACLAuthenticatedRead:      true,
		s3.ObjectCannedACLAwsExecRead:            true,
		s3.ObjectCannedACLBucketOwnerRead:        true,
		s3.ObjectCannedACLBucketOwnerFullControl: true,
	}

	sentenceJoin := func(m map[string]bool) string {
		keys := make([]string, 0, len(m))
		for k := range m {
			keys = append(keys, fmt.Sprintf("%q", k))
		}
		sort.Strings(keys)

		length := len(keys)
		words := make([]string, length)
		copy(words, keys)

		words[length-1] = fmt.Sprintf("or %s", words[length-1])
		return strings.Join(words, ", ")
	}

	if _, ok := cannedAcls[value]; !ok {
		errors = append(errors, fmt.Errorf(
			"%q contains an invalid canned ACL type %q. Valid types are either %s",
			k, value, sentenceJoin(cannedAcls)))
	}
	return
}

func validateS3BucketObjectStorageClassType(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)

	storageClass := map[string]bool{
		s3.StorageClassStandard:          true,
		s3.StorageClassReducedRedundancy: true,
		s3.StorageClassStandardIa:        true,
	}

	if _, ok := storageClass[value]; !ok {
		errors = append(errors, fmt.Errorf(
			"%q contains an invalid Storage Class type %q. Valid types are either %q, %q, or %q",
			k, value, s3.StorageClassStandard, s3.StorageClassReducedRedundancy,
			s3.StorageClassStandardIa))
	}
	return
}

func validateS3BucketObjectServerSideEncryption(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)

	serverSideEncryption := map[string]bool{
		s3.ServerSideEncryptionAes256: true,
		s3.ServerSideEncryptionAwsKms: true,
	}

	if _, ok := serverSideEncryption[value]; !ok {
		errors = append(errors, fmt.Errorf(
			"%q contains an invalid Server Side Encryption value %q. Valid values are %q and %q",
			k, value, s3.ServerSideEncryptionAes256, s3.ServerSideEncryptionAwsKms))
	}
	return
}
<commit_msg>add support for govcloud s3 object tags<commit_after>package aws

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"sort"
	"strings"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/go-homedir"

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

func resourceAwsS3BucketObject() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsS3BucketObjectPut,
		Read:   resourceAwsS3BucketObjectRead,
		Update: resourceAwsS3BucketObjectPut,
		Delete: resourceAwsS3BucketObjectDelete,

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

			"acl": {
				Type:         schema.TypeString,
				Default:      "private",
				Optional:     true,
				ValidateFunc: validateS3BucketObjectAclType,
			},

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

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

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

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

			"content_type": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"key": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"source": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"content"},
			},

			"content": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"source"},
			},

			"storage_class": {
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				ValidateFunc: validateS3BucketObjectStorageClassType,
			},

			"server_side_encryption": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateS3BucketObjectServerSideEncryption,
				Computed:     true,
			},

			"kms_key_id": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateArn,
			},

			"etag": {
				Type: schema.TypeString,
				// This will conflict with SSE-C and SSE-KMS encryption and multi-part upload
				// if/when it's actually implemented. The Etag then won't match raw-file MD5.
				// See http://docs.aws.amazon.com/AmazonS3/latest/API/RESTCommonResponseHeaders.html
				Optional:      true,
				Computed:      true,
				ConflictsWith: []string{"kms_key_id", "server_side_encryption"},
			},

			"version_id": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),

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

func resourceAwsS3BucketObjectPut(d *schema.ResourceData, meta interface{}) error {
	s3conn := meta.(*AWSClient).s3conn

	restricted := meta.(*AWSClient).IsChinaCloud()

	var body io.ReadSeeker

	if v, ok := d.GetOk("source"); ok {
		source := v.(string)
		path, err := homedir.Expand(source)
		if err != nil {
			return fmt.Errorf("Error expanding homedir in source (%s): %s", source, err)
		}
		file, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Error opening S3 bucket object source (%s): %s", source, err)
		}

		body = file
	} else if v, ok := d.GetOk("content"); ok {
		content := v.(string)
		body = bytes.NewReader([]byte(content))
	} else {
		return fmt.Errorf("Must specify \"source\" or \"content\" field")
	}

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

	putInput := &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
		ACL:    aws.String(d.Get("acl").(string)),
		Body:   body,
	}

	if v, ok := d.GetOk("storage_class"); ok {
		putInput.StorageClass = aws.String(v.(string))
	}

	if v, ok := d.GetOk("cache_control"); ok {
		putInput.CacheControl = aws.String(v.(string))
	}

	if v, ok := d.GetOk("content_type"); ok {
		putInput.ContentType = aws.String(v.(string))
	}

	if v, ok := d.GetOk("content_encoding"); ok {
		putInput.ContentEncoding = aws.String(v.(string))
	}

	if v, ok := d.GetOk("content_language"); ok {
		putInput.ContentLanguage = aws.String(v.(string))
	}

	if v, ok := d.GetOk("content_disposition"); ok {
		putInput.ContentDisposition = aws.String(v.(string))
	}

	if v, ok := d.GetOk("server_side_encryption"); ok {
		putInput.ServerSideEncryption = aws.String(v.(string))
	}

	if v, ok := d.GetOk("kms_key_id"); ok {
		putInput.SSEKMSKeyId = aws.String(v.(string))
		putInput.ServerSideEncryption = aws.String(s3.ServerSideEncryptionAwsKms)
	}

	if v, ok := d.GetOk("tags"); ok {
		if restricted {
			return fmt.Errorf("This region does not allow for tags on S3 objects")
		}

		// The tag-set must be encoded as URL Query parameters.
		values := url.Values{}
		for k, v := range v.(map[string]interface{}) {
			values.Add(k, v.(string))
		}
		putInput.Tagging = aws.String(values.Encode())
	}

	if v, ok := d.GetOk("website_redirect"); ok {
		putInput.WebsiteRedirectLocation = aws.String(v.(string))
	}

	resp, err := s3conn.PutObject(putInput)
	if err != nil {
		return fmt.Errorf("Error putting object in S3 bucket (%s): %s", bucket, err)
	}

	// See https://forums.aws.amazon.com/thread.jspa?threadID=44003
	d.Set("etag", strings.Trim(*resp.ETag, `"`))

	d.Set("version_id", resp.VersionId)
	d.SetId(key)
	return resourceAwsS3BucketObjectRead(d, meta)
}

func resourceAwsS3BucketObjectRead(d *schema.ResourceData, meta interface{}) error {
	s3conn := meta.(*AWSClient).s3conn

	restricted := meta.(*AWSClient).IsChinaCloud()

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

	resp, err := s3conn.HeadObject(
		&s3.HeadObjectInput{
			Bucket: aws.String(bucket),
			Key:    aws.String(key),
		})

	if err != nil {
		// If S3 returns a 404 Request Failure, mark the object as destroyed
		if awsErr, ok := err.(awserr.RequestFailure); ok && awsErr.StatusCode() == 404 {
			d.SetId("")
			log.Printf("[WARN] Error Reading Object (%s), object not found (HTTP status 404)", key)
			return nil
		}
		return err
	}
	log.Printf("[DEBUG] Reading S3 Bucket Object meta: %s", resp)

	d.Set("cache_control", resp.CacheControl)
	d.Set("content_disposition", resp.ContentDisposition)
	d.Set("content_encoding", resp.ContentEncoding)
	d.Set("content_language", resp.ContentLanguage)
	d.Set("content_type", resp.ContentType)
	d.Set("version_id", resp.VersionId)
	d.Set("server_side_encryption", resp.ServerSideEncryption)
	d.Set("website_redirect", resp.WebsiteRedirectLocation)

	// Only set non-default KMS key ID (one that doesn't match default)
	if resp.SSEKMSKeyId != nil {
		// retrieve S3 KMS Default Master Key
		kmsconn := meta.(*AWSClient).kmsconn
		kmsresp, err := kmsconn.DescribeKey(&kms.DescribeKeyInput{
			KeyId: aws.String("alias/aws/s3"),
		})
		if err != nil {
			return fmt.Errorf("Failed to describe default S3 KMS key (alias/aws/s3): %s", err)
		}

		if *resp.SSEKMSKeyId != *kmsresp.KeyMetadata.Arn {
			log.Printf("[DEBUG] S3 object is encrypted using a non-default KMS Key ID: %s", *resp.SSEKMSKeyId)
			d.Set("kms_key_id", resp.SSEKMSKeyId)
		}
	}
	d.Set("etag", strings.Trim(*resp.ETag, `"`))

	// The "STANDARD" (which is also the default) storage
	// class when set would not be included in the results.
	d.Set("storage_class", s3.StorageClassStandard)
	if resp.StorageClass != nil {
		d.Set("storage_class", resp.StorageClass)
	}

	if !restricted {
		tagResp, err := s3conn.GetObjectTagging(
			&s3.GetObjectTaggingInput{
				Bucket: aws.String(bucket),
				Key:    aws.String(key),
			})
		if err != nil {
			return fmt.Errorf("Failed to get object tags (bucket: %s, key: %s): %s", bucket, key, err)
		}
		d.Set("tags", tagsToMapS3(tagResp.TagSet))
	}

	return nil
}

func resourceAwsS3BucketObjectDelete(d *schema.ResourceData, meta interface{}) error {
	s3conn := meta.(*AWSClient).s3conn

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

	if _, ok := d.GetOk("version_id"); ok {
		// Bucket is versioned, we need to delete all versions
		vInput := s3.ListObjectVersionsInput{
			Bucket: aws.String(bucket),
			Prefix: aws.String(key),
		}
		out, err := s3conn.ListObjectVersions(&vInput)
		if err != nil {
			return fmt.Errorf("Failed listing S3 object versions: %s", err)
		}

		for _, v := range out.Versions {
			input := s3.DeleteObjectInput{
				Bucket:    aws.String(bucket),
				Key:       aws.String(key),
				VersionId: v.VersionId,
			}
			_, err := s3conn.DeleteObject(&input)
			if err != nil {
				return fmt.Errorf("Error deleting S3 object version of %s:\n %s:\n %s",
					key, v, err)
			}
		}
	} else {
		// Just delete the object
		input := s3.DeleteObjectInput{
			Bucket: aws.String(bucket),
			Key:    aws.String(key),
		}
		_, err := s3conn.DeleteObject(&input)
		if err != nil {
			return fmt.Errorf("Error deleting S3 bucket object: %s  Bucket: %q Object: %q", err, bucket, key)
		}
	}

	return nil
}

func validateS3BucketObjectAclType(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)

	cannedAcls := map[string]bool{
		s3.ObjectCannedACLPrivate:                true,
		s3.ObjectCannedACLPublicRead:             true,
		s3.ObjectCannedACLPublicReadWrite:        true,
		s3.ObjectCannedACLAuthenticatedRead:      true,
		s3.ObjectCannedACLAwsExecRead:            true,
		s3.ObjectCannedACLBucketOwnerRead:        true,
		s3.ObjectCannedACLBucketOwnerFullControl: true,
	}

	sentenceJoin := func(m map[string]bool) string {
		keys := make([]string, 0, len(m))
		for k := range m {
			keys = append(keys, fmt.Sprintf("%q", k))
		}
		sort.Strings(keys)

		length := len(keys)
		words := make([]string, length)
		copy(words, keys)

		words[length-1] = fmt.Sprintf("or %s", words[length-1])
		return strings.Join(words, ", ")
	}

	if _, ok := cannedAcls[value]; !ok {
		errors = append(errors, fmt.Errorf(
			"%q contains an invalid canned ACL type %q. Valid types are either %s",
			k, value, sentenceJoin(cannedAcls)))
	}
	return
}

func validateS3BucketObjectStorageClassType(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)

	storageClass := map[string]bool{
		s3.StorageClassStandard:          true,
		s3.StorageClassReducedRedundancy: true,
		s3.StorageClassStandardIa:        true,
	}

	if _, ok := storageClass[value]; !ok {
		errors = append(errors, fmt.Errorf(
			"%q contains an invalid Storage Class type %q. Valid types are either %q, %q, or %q",
			k, value, s3.StorageClassStandard, s3.StorageClassReducedRedundancy,
			s3.StorageClassStandardIa))
	}
	return
}

func validateS3BucketObjectServerSideEncryption(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)

	serverSideEncryption := map[string]bool{
		s3.ServerSideEncryptionAes256: true,
		s3.ServerSideEncryptionAwsKms: true,
	}

	if _, ok := serverSideEncryption[value]; !ok {
		errors = append(errors, fmt.Errorf(
			"%q contains an invalid Server Side Encryption value %q. Valid values are %q and %q",
			k, value, s3.ServerSideEncryptionAes256, s3.ServerSideEncryptionAwsKms))
	}
	return
}
<|endoftext|>


<commit_before>package harvest

import (
	"fmt"
	"time"
)

type TaskAssignmentResponse struct {
	TaskAssignment *TaskAssignment `json:"task_assignment"`
}

type TaskAssignment struct {
	ID          int64     `json:"id"`
	ProjectID   int64     `json:"project_id"`
	TaskID      int64     `json:"task_id"`
	Billable    bool      `json:"billable"`
	Deactivated bool      `json:"deactivated"`
	Budget      *float64  `json:"budget"`
	HourlyRate  *float64  `json:"hourly_rate"`
	UpdatedAt   time.Time `json:"updated_at"`
	CreatedAt   time.Time `json:"created_at"`
}

func (a *API) GetTaskAssignments(projectID int64, args Arguments) (taskassignments []*TaskAssignment, err error) {
	taskAssignmentsResponse := make([]*TaskAssignmentResponse, 0)
	path := fmt.Sprintf("/projects/%v/task_assignments", projectID)
	err = a.Get(path, args, &taskAssignmentsResponse)
	for _, ta := range taskAssignmentsResponse {
		taskassignments = append(taskassignments, ta.TaskAssignment)
	}
	return taskassignments, err
}

func (a *API) GetTaskAssignment(projectID int64, taskAssignmentID int64, args Arguments) (taskassignment *TaskAssignment, err error) {
	taskAssignmentResponse := TaskAssignmentResponse{}
	path := fmt.Sprintf("/projects/%v/task_assignments/%v", projectID, taskAssignmentID)
	err = a.Get(path, args, &taskAssignmentResponse)
	return taskAssignmentResponse.TaskAssignment, err
}
<commit_msg>Adds missing TaskAssignment.Estimate.<commit_after>package harvest

import (
	"fmt"
	"time"
)

type TaskAssignmentResponse struct {
	TaskAssignment *TaskAssignment `json:"task_assignment"`
}

type TaskAssignment struct {
	ID          int64     `json:"id"`
	ProjectID   int64     `json:"project_id"`
	TaskID      int64     `json:"task_id"`
	Billable    bool      `json:"billable"`
	Deactivated bool      `json:"deactivated"`
	Budget      *float64  `json:"budget"`
	HourlyRate  *float64  `json:"hourly_rate"`
	Estimate    *float64  `json:"estimate"`
	UpdatedAt   time.Time `json:"updated_at"`
	CreatedAt   time.Time `json:"created_at"`
}

func (a *API) GetTaskAssignments(projectID int64, args Arguments) (taskassignments []*TaskAssignment, err error) {
	taskAssignmentsResponse := make([]*TaskAssignmentResponse, 0)
	path := fmt.Sprintf("/projects/%v/task_assignments", projectID)
	err = a.Get(path, args, &taskAssignmentsResponse)
	for _, ta := range taskAssignmentsResponse {
		taskassignments = append(taskassignments, ta.TaskAssignment)
	}
	return taskassignments, err
}

func (a *API) GetTaskAssignment(projectID int64, taskAssignmentID int64, args Arguments) (taskassignment *TaskAssignment, err error) {
	taskAssignmentResponse := TaskAssignmentResponse{}
	path := fmt.Sprintf("/projects/%v/task_assignments/%v", projectID, taskAssignmentID)
	err = a.Get(path, args, &taskAssignmentResponse)
	return taskAssignmentResponse.TaskAssignment, err
}
<|endoftext|>


<commit_before>package models

import (
	"errors"
	"fmt"
	"socialapi/config"
	"time"

	"github.com/VerbalExpressions/GoVerbalExpressions"
	"github.com/jinzhu/gorm"
	"github.com/koding/bongo"
)

var mentionRegex = verbalexpressions.New().
	Find("@").
	BeginCapture().
	Word().
	EndCapture().
	Regex()

type ChannelMessage struct {
	// unique identifier of the channel message
	Id int64 `json:"id,string"`

	// Body of the mesage
	Body string `json:"body"`

	// Generated Slug for body
	Slug string `json:"slug"                               sql:"NOT NULL;TYPE:VARCHAR(100);"`

	// type of the message
	TypeConstant string `json:"typeConstant"               sql:"NOT NULL;TYPE:VARCHAR(100);"`

	// Creator of the channel message
	AccountId int64 `json:"accountId,string"               sql:"NOT NULL"`

	// in which channel this message is created
	InitialChannelId int64 `json:"initialChannelId,string" sql:"NOT NULL"`

	// holds troll, unsafe, etc
	MetaBits int16 `json:"-"`

	// Creation date of the message
	CreatedAt time.Time `json:"createdAt"                  sql:"DEFAULT:CURRENT_TIMESTAMP"`

	// Modification date of the message
	UpdatedAt time.Time `json:"updatedAt"                  sql:"DEFAULT:CURRENT_TIMESTAMP"`

	// Deletion date of the channel message
	DeletedAt time.Time `json:"deletedAt"`

	// Extra data storage
	Payload gorm.Hstore `json:"payload,omitempty"`
}

func (c *ChannelMessage) BeforeCreate() {
	if res, err := c.isExemptContent(); err == nil && res {
		c.MetaBits = updateTrollModeBit(c.MetaBits)
	}

	c.DeletedAt = ZeroDate()
}

func (c *ChannelMessage) BeforeUpdate() {
	if res, err := c.isExemptContent(); err == nil && res {
		c.MetaBits = updateTrollModeBit(c.MetaBits)
	}

	c.DeletedAt = ZeroDate()
}

func (c *ChannelMessage) AfterCreate() {
	bongo.B.AfterCreate(c)
}

func (c *ChannelMessage) AfterUpdate() {
	bongo.B.AfterUpdate(c)
}

func (c ChannelMessage) AfterDelete() {
	bongo.B.AfterDelete(c)
}

func (c ChannelMessage) GetId() int64 {
	return c.Id
}

func (c ChannelMessage) TableName() string {
	return "api.channel_message"
}

const (
	ChannelMessage_TYPE_POST            = "post"
	ChannelMessage_TYPE_REPLY           = "reply"
	ChannelMessage_TYPE_JOIN            = "join"
	ChannelMessage_TYPE_LEAVE           = "leave"
	ChannelMessage_TYPE_CHAT            = "chat"
	ChannelMessage_TYPE_PRIVATE_MESSAGE = "privatemessage"
)

func NewChannelMessage() *ChannelMessage {
	return &ChannelMessage{}
}

func (c *ChannelMessage) ById(id int64) error {
	return bongo.B.ById(c, id)
}

func (c *ChannelMessage) One(q *bongo.Query) error {
	return bongo.B.One(c, c, q)
}

func (c *ChannelMessage) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(c, data, q)
}

func (c *ChannelMessage) UpdateMulti(rest ...map[string]interface{}) error {
	return bongo.B.UpdateMulti(c, rest...)
}

func (c *ChannelMessage) CountWithQuery(q *bongo.Query) (int, error) {
	return bongo.B.CountWithQuery(c, q)
}

func (c *ChannelMessage) isExemptContent() (bool, error) {
	// set meta bits if only message is post or a reply
	if c.TypeConstant != ChannelMessage_TYPE_POST &&
		c.TypeConstant != ChannelMessage_TYPE_REPLY {
		return false, nil
	}

	if c.AccountId == 0 && c.Id != 0 {
		if err := c.ById(c.Id); err != nil {
			return false, err
		}
	} else {
		return false, fmt.Errorf("Couldnt find accountId from content %+v", c)
	}

	account, err := FetchAccountFromCache(c.AccountId)
	if err != nil {
		return false, err
	}

	if account == nil {
		return false, fmt.Errorf("Account is nil, accountId:%d", c.AccountId)
	}

	if account.IsTroll {
		return true, nil
	}

	return false, nil
}

func bodyLenCheck(body string) error {
	if len(body) < config.Get().Limits.MessageBodyMinLen {
		return fmt.Errorf("Message Body Length should be greater than %d, yours is %d ", config.Get().Limits.MessageBodyMinLen, len(body))
	}

	return nil
}

// todo create a new message while updating the channel_message and delete other
// cases, since deletion is a soft delete, old instances will still be there
func (c *ChannelMessage) Update() error {
	if err := bodyLenCheck(c.Body); err != nil {
		return err
	}
	// only update body
	err := bongo.B.UpdatePartial(c,
		map[string]interface{}{
			"body": c.Body,
		},
	)
	return err
}

func (c *ChannelMessage) Create() error {
	if err := bodyLenCheck(c.Body); err != nil {
		return err
	}

	var err error
	c, err = Slugify(c)
	if err != nil {
		return err
	}

	return bongo.B.Create(c)
}

// CreateRaw creates a new channel message without effected by auto generated createdAt
// and updatedAt values
func (c *ChannelMessage) CreateRaw() error {
	insertSql := "INSERT INTO " +
		c.TableName() +
		` ("body","slug","type_constant","account_id","initial_channel_id",` +
		`"created_at","updated_at","deleted_at","payload") ` +
		"VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9) " +
		"RETURNING ID"

	return bongo.B.DB.CommonDB().QueryRow(insertSql, c.Body, c.Slug, c.TypeConstant, c.AccountId, c.InitialChannelId,
		c.CreatedAt, c.UpdatedAt, c.DeletedAt, c.Payload).Scan(&c.Id)
}

func (c *ChannelMessage) Delete() error {
	return bongo.B.Delete(c)
}

func (c *ChannelMessage) FetchByIds(ids []int64) ([]ChannelMessage, error) {
	var messages []ChannelMessage

	if len(ids) == 0 {
		return messages, nil
	}

	if err := bongo.B.FetchByIds(c, &messages, ids); err != nil {
		return nil, err
	}

	return messages, nil
}

func (c *ChannelMessage) BuildMessages(query *Query, messages []ChannelMessage) ([]*ChannelMessageContainer, error) {
	containers := make([]*ChannelMessageContainer, len(messages))
	if len(containers) == 0 {
		return containers, nil
	}

	for i, message := range messages {
		d := NewChannelMessage()
		*d = message
		data, err := d.BuildMessage(query)
		if err != nil {
			return containers, err
		}
		containers[i] = data
	}

	return containers, nil
}

func (c *ChannelMessage) BuildMessage(query *Query) (*ChannelMessageContainer, error) {
	cmc, err := c.FetchRelatives(query)
	if err != nil {
		return nil, err
	}

	mr := NewMessageReply()
	mr.MessageId = c.Id
	q := query
	q.Limit = 3
	replies, err := mr.List(query)
	if err != nil {
		return nil, err
	}

	repliesCount, err := mr.Count()
	if err != nil {
		return nil, err
	}
	cmc.RepliesCount = repliesCount

	cmc.IsFollowed, err = c.CheckIsMessageFollowed(query)
	if err != nil {
		return nil, err
	}

	populatedChannelMessagesReplies := make([]*ChannelMessageContainer, len(replies))
	for rl := 0; rl < len(replies); rl++ {
		cmrc, err := replies[rl].FetchRelatives(query)
		if err != nil {
			return nil, err
		}
		populatedChannelMessagesReplies[rl] = cmrc
	}

	cmc.Replies = populatedChannelMessagesReplies
	return cmc, nil
}

func (c *ChannelMessage) CheckIsMessageFollowed(query *Query) (bool, error) {
	channel := NewChannel()
	if err := channel.FetchPinnedActivityChannel(query.AccountId, query.GroupName); err != nil {
		if err == gorm.RecordNotFound {
			return false, nil
		}
		return false, err
	}

	cml := NewChannelMessageList()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channel.Id,
			"message_id": c.Id,
		},
	}
	if err := cml.One(q); err != nil {
		if err == gorm.RecordNotFound {
			return false, nil
		}

		return false, err
	}

	return true, nil
}

func (c *ChannelMessage) BuildEmptyMessageContainer() (*ChannelMessageContainer, error) {
	if c.Id == 0 {
		return nil, errors.New("Channel message id is not set")
	}
	container := NewChannelMessageContainer()
	container.Message = c

	oldId, err := FetchAccountOldIdByIdFromCache(c.AccountId)
	if err != nil {
		return nil, err
	}

	container.AccountOldId = oldId

	interactionContainer := NewInteractionContainer()
	interactionContainer.ActorsPreview = make([]string, 0)
	interactionContainer.IsInteracted = false
	interactionContainer.ActorsCount = 0

	container.Interactions = make(map[string]*InteractionContainer)
	container.Interactions["like"] = interactionContainer

	return container, nil
}

func (c *ChannelMessage) FetchRelatives(query *Query) (*ChannelMessageContainer, error) {
	container, err := c.BuildEmptyMessageContainer()
	if err != nil {
		return nil, err
	}

	i := NewInteraction()
	i.MessageId = c.Id

	// get preview
	query.Type = "like"
	query.Limit = 3
	interactorIds, err := i.List(query)
	if err != nil {
		return nil, err
	}

	oldIds, err := FetchOldIdsByAccountIds(interactorIds)
	if err != nil {
		return nil, err
	}

	interactionContainer := NewInteractionContainer()
	interactionContainer.ActorsPreview = oldIds

	// check if the current user is interacted in this thread
	isInteracted, err := i.IsInteracted(query.AccountId)
	if err != nil {
		return nil, err
	}

	interactionContainer.IsInteracted = isInteracted

	// fetch interaction count
	count, err := i.Count(query.Type)
	if err != nil {
		return nil, err
	}

	interactionContainer.ActorsCount = count

	container.Interactions["like"] = interactionContainer
	return container, nil
}

func generateMessageListQuery(channelId int64, q *Query) *bongo.Query {
	messageType := q.Type
	if messageType == "" {
		messageType = ChannelMessage_TYPE_POST
	}

	return &bongo.Query{
		Selector: map[string]interface{}{
			"account_id":         q.AccountId,
			"initial_channel_id": channelId,
			"type_constant":      messageType,
		},
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort: map[string]string{
			"created_at": "DESC",
		},
	}
}

func (c *ChannelMessage) FetchMessagesByChannelId(channelId int64, q *Query) ([]ChannelMessage, error) {
	query := generateMessageListQuery(channelId, q)

	var messages []ChannelMessage
	if err := c.Some(&messages, query); err != nil {
		return nil, err
	}

	if messages == nil {
		return make([]ChannelMessage, 0), nil
	}
	return messages, nil
}

func (c *ChannelMessage) GetMentionedUsernames() []string {
	flattened := make([]string, 0)

	res := mentionRegex.FindAllStringSubmatch(c.Body, -1)
	if len(res) == 0 {
		return flattened
	}

	participants := map[string]struct{}{}
	// remove duplicate mentions
	for _, ele := range res {
		participants[ele[1]] = struct{}{}
	}

	for participant := range participants {
		flattened = append(flattened, participant)
	}

	return flattened
}

func (c *ChannelMessage) FetchTotalMessageCount(q *Query) (int, error) {
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id":    q.AccountId,
			"type_constant": q.Type,
		},
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
	}

	return c.CountWithQuery(query)
}

func (c *ChannelMessage) FetchMessageIds(q *Query) ([]int64, error) {
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id":    q.AccountId,
			"type_constant": q.Type,
		},
		Pluck:      "id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort: map[string]string{
			"created_at": "DESC",
		},
	}

	var messageIds []int64
	if err := c.Some(&messageIds, query); err != nil {
		return nil, err
	}

	if messageIds == nil {
		return make([]int64, 0), nil
	}

	return messageIds, nil
}

// BySlug fetchs channel message by its slug
// checks if message is in the channel or not
func (c *ChannelMessage) BySlug(query *Query) error {
	if query.Slug == "" {
		return errors.New("slug is not set")
	}

	// fetch message itself
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"slug": query.Slug,
		},
	}

	if err := c.One(q); err != nil {
		return nil
	}

	// fetch channel by group name
	query.Name = query.GroupName
	query.Type = Channel_TYPE_GROUP
	ch := NewChannel()
	channel, err := ch.ByName(query)
	if err != nil {
		return err
	}

	if channel.Id == 0 {
		return errors.New("channel is not found")
	}

	// check if message is in the channel
	cml := NewChannelMessageList()
	res, err := cml.IsInChannel(c.Id, channel.Id)
	if err != nil {
		return err
	}

	// if message is not in the channel
	if !res {
		return gorm.RecordNotFound
	}

	return nil
}
<commit_msg>Migration: Add body update method (bypass callbacks)<commit_after>package models

import (
	"errors"
	"fmt"
	"socialapi/config"
	"time"

	"github.com/VerbalExpressions/GoVerbalExpressions"
	"github.com/jinzhu/gorm"
	"github.com/koding/bongo"
)

var mentionRegex = verbalexpressions.New().
	Find("@").
	BeginCapture().
	Word().
	EndCapture().
	Regex()

type ChannelMessage struct {
	// unique identifier of the channel message
	Id int64 `json:"id,string"`

	// Body of the mesage
	Body string `json:"body"`

	// Generated Slug for body
	Slug string `json:"slug"                               sql:"NOT NULL;TYPE:VARCHAR(100);"`

	// type of the message
	TypeConstant string `json:"typeConstant"               sql:"NOT NULL;TYPE:VARCHAR(100);"`

	// Creator of the channel message
	AccountId int64 `json:"accountId,string"               sql:"NOT NULL"`

	// in which channel this message is created
	InitialChannelId int64 `json:"initialChannelId,string" sql:"NOT NULL"`

	// holds troll, unsafe, etc
	MetaBits int16 `json:"-"`

	// Creation date of the message
	CreatedAt time.Time `json:"createdAt"                  sql:"DEFAULT:CURRENT_TIMESTAMP"`

	// Modification date of the message
	UpdatedAt time.Time `json:"updatedAt"                  sql:"DEFAULT:CURRENT_TIMESTAMP"`

	// Deletion date of the channel message
	DeletedAt time.Time `json:"deletedAt"`

	// Extra data storage
	Payload gorm.Hstore `json:"payload,omitempty"`
}

func (c *ChannelMessage) BeforeCreate() {
	if res, err := c.isExemptContent(); err == nil && res {
		c.MetaBits = updateTrollModeBit(c.MetaBits)
	}

	c.DeletedAt = ZeroDate()
}

func (c *ChannelMessage) BeforeUpdate() {
	if res, err := c.isExemptContent(); err == nil && res {
		c.MetaBits = updateTrollModeBit(c.MetaBits)
	}

	c.DeletedAt = ZeroDate()
}

func (c *ChannelMessage) AfterCreate() {
	bongo.B.AfterCreate(c)
}

func (c *ChannelMessage) AfterUpdate() {
	bongo.B.AfterUpdate(c)
}

func (c ChannelMessage) AfterDelete() {
	bongo.B.AfterDelete(c)
}

func (c ChannelMessage) GetId() int64 {
	return c.Id
}

func (c ChannelMessage) TableName() string {
	return "api.channel_message"
}

const (
	ChannelMessage_TYPE_POST            = "post"
	ChannelMessage_TYPE_REPLY           = "reply"
	ChannelMessage_TYPE_JOIN            = "join"
	ChannelMessage_TYPE_LEAVE           = "leave"
	ChannelMessage_TYPE_CHAT            = "chat"
	ChannelMessage_TYPE_PRIVATE_MESSAGE = "privatemessage"
)

func NewChannelMessage() *ChannelMessage {
	return &ChannelMessage{}
}

func (c *ChannelMessage) ById(id int64) error {
	return bongo.B.ById(c, id)
}

func (c *ChannelMessage) One(q *bongo.Query) error {
	return bongo.B.One(c, c, q)
}

func (c *ChannelMessage) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(c, data, q)
}

func (c *ChannelMessage) UpdateMulti(rest ...map[string]interface{}) error {
	return bongo.B.UpdateMulti(c, rest...)
}

func (c *ChannelMessage) CountWithQuery(q *bongo.Query) (int, error) {
	return bongo.B.CountWithQuery(c, q)
}

func (c *ChannelMessage) isExemptContent() (bool, error) {
	// set meta bits if only message is post or a reply
	if c.TypeConstant != ChannelMessage_TYPE_POST &&
		c.TypeConstant != ChannelMessage_TYPE_REPLY {
		return false, nil
	}

	if c.AccountId == 0 && c.Id != 0 {
		if err := c.ById(c.Id); err != nil {
			return false, err
		}
	} else {
		return false, fmt.Errorf("Couldnt find accountId from content %+v", c)
	}

	account, err := FetchAccountFromCache(c.AccountId)
	if err != nil {
		return false, err
	}

	if account == nil {
		return false, fmt.Errorf("Account is nil, accountId:%d", c.AccountId)
	}

	if account.IsTroll {
		return true, nil
	}

	return false, nil
}

func bodyLenCheck(body string) error {
	if len(body) < config.Get().Limits.MessageBodyMinLen {
		return fmt.Errorf("Message Body Length should be greater than %d, yours is %d ", config.Get().Limits.MessageBodyMinLen, len(body))
	}

	return nil
}

// todo create a new message while updating the channel_message and delete other
// cases, since deletion is a soft delete, old instances will still be there
func (c *ChannelMessage) Update() error {
	if err := bodyLenCheck(c.Body); err != nil {
		return err
	}
	// only update body
	err := bongo.B.UpdatePartial(c,
		map[string]interface{}{
			"body": c.Body,
		},
	)
	return err
}

func (c *ChannelMessage) Create() error {
	if err := bodyLenCheck(c.Body); err != nil {
		return err
	}

	var err error
	c, err = Slugify(c)
	if err != nil {
		return err
	}

	return bongo.B.Create(c)
}

// CreateRaw creates a new channel message without effected by auto generated createdAt
// and updatedAt values
func (c *ChannelMessage) CreateRaw() error {
	insertSql := "INSERT INTO " +
		c.TableName() +
		` ("body","slug","type_constant","account_id","initial_channel_id",` +
		`"created_at","updated_at","deleted_at","payload") ` +
		"VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9) " +
		"RETURNING ID"

	return bongo.B.DB.CommonDB().QueryRow(insertSql, c.Body, c.Slug, c.TypeConstant, c.AccountId, c.InitialChannelId,
		c.CreatedAt, c.UpdatedAt, c.DeletedAt, c.Payload).Scan(&c.Id)
}

// UpdateBodyRaw updates message body without effecting createdAt/UpdatedAt
// timestamps
func (c *ChannelMessage) UpdateBodyRaw() error {
	updateSql := fmt.Sprintf("UPDATE %s SET body=? WHERE id=?", c.TableName())

	return bongo.B.DB.Exec(updateSql, c.Body, c.Id).Error
}

func (c *ChannelMessage) Delete() error {
	return bongo.B.Delete(c)
}

func (c *ChannelMessage) FetchByIds(ids []int64) ([]ChannelMessage, error) {
	var messages []ChannelMessage

	if len(ids) == 0 {
		return messages, nil
	}

	if err := bongo.B.FetchByIds(c, &messages, ids); err != nil {
		return nil, err
	}

	return messages, nil
}

func (c *ChannelMessage) BuildMessages(query *Query, messages []ChannelMessage) ([]*ChannelMessageContainer, error) {
	containers := make([]*ChannelMessageContainer, len(messages))
	if len(containers) == 0 {
		return containers, nil
	}

	for i, message := range messages {
		d := NewChannelMessage()
		*d = message
		data, err := d.BuildMessage(query)
		if err != nil {
			return containers, err
		}
		containers[i] = data
	}

	return containers, nil
}

func (c *ChannelMessage) BuildMessage(query *Query) (*ChannelMessageContainer, error) {
	cmc, err := c.FetchRelatives(query)
	if err != nil {
		return nil, err
	}

	mr := NewMessageReply()
	mr.MessageId = c.Id
	q := query
	q.Limit = 3
	replies, err := mr.List(query)
	if err != nil {
		return nil, err
	}

	repliesCount, err := mr.Count()
	if err != nil {
		return nil, err
	}
	cmc.RepliesCount = repliesCount

	cmc.IsFollowed, err = c.CheckIsMessageFollowed(query)
	if err != nil {
		return nil, err
	}

	populatedChannelMessagesReplies := make([]*ChannelMessageContainer, len(replies))
	for rl := 0; rl < len(replies); rl++ {
		cmrc, err := replies[rl].FetchRelatives(query)
		if err != nil {
			return nil, err
		}
		populatedChannelMessagesReplies[rl] = cmrc
	}

	cmc.Replies = populatedChannelMessagesReplies
	return cmc, nil
}

func (c *ChannelMessage) CheckIsMessageFollowed(query *Query) (bool, error) {
	channel := NewChannel()
	if err := channel.FetchPinnedActivityChannel(query.AccountId, query.GroupName); err != nil {
		if err == gorm.RecordNotFound {
			return false, nil
		}
		return false, err
	}

	cml := NewChannelMessageList()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channel.Id,
			"message_id": c.Id,
		},
	}
	if err := cml.One(q); err != nil {
		if err == gorm.RecordNotFound {
			return false, nil
		}

		return false, err
	}

	return true, nil
}

func (c *ChannelMessage) BuildEmptyMessageContainer() (*ChannelMessageContainer, error) {
	if c.Id == 0 {
		return nil, errors.New("Channel message id is not set")
	}
	container := NewChannelMessageContainer()
	container.Message = c

	oldId, err := FetchAccountOldIdByIdFromCache(c.AccountId)
	if err != nil {
		return nil, err
	}

	container.AccountOldId = oldId

	interactionContainer := NewInteractionContainer()
	interactionContainer.ActorsPreview = make([]string, 0)
	interactionContainer.IsInteracted = false
	interactionContainer.ActorsCount = 0

	container.Interactions = make(map[string]*InteractionContainer)
	container.Interactions["like"] = interactionContainer

	return container, nil
}

func (c *ChannelMessage) FetchRelatives(query *Query) (*ChannelMessageContainer, error) {
	container, err := c.BuildEmptyMessageContainer()
	if err != nil {
		return nil, err
	}

	i := NewInteraction()
	i.MessageId = c.Id

	// get preview
	query.Type = "like"
	query.Limit = 3
	interactorIds, err := i.List(query)
	if err != nil {
		return nil, err
	}

	oldIds, err := FetchOldIdsByAccountIds(interactorIds)
	if err != nil {
		return nil, err
	}

	interactionContainer := NewInteractionContainer()
	interactionContainer.ActorsPreview = oldIds

	// check if the current user is interacted in this thread
	isInteracted, err := i.IsInteracted(query.AccountId)
	if err != nil {
		return nil, err
	}

	interactionContainer.IsInteracted = isInteracted

	// fetch interaction count
	count, err := i.Count(query.Type)
	if err != nil {
		return nil, err
	}

	interactionContainer.ActorsCount = count

	container.Interactions["like"] = interactionContainer
	return container, nil
}

func generateMessageListQuery(channelId int64, q *Query) *bongo.Query {
	messageType := q.Type
	if messageType == "" {
		messageType = ChannelMessage_TYPE_POST
	}

	return &bongo.Query{
		Selector: map[string]interface{}{
			"account_id":         q.AccountId,
			"initial_channel_id": channelId,
			"type_constant":      messageType,
		},
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort: map[string]string{
			"created_at": "DESC",
		},
	}
}

func (c *ChannelMessage) FetchMessagesByChannelId(channelId int64, q *Query) ([]ChannelMessage, error) {
	query := generateMessageListQuery(channelId, q)

	var messages []ChannelMessage
	if err := c.Some(&messages, query); err != nil {
		return nil, err
	}

	if messages == nil {
		return make([]ChannelMessage, 0), nil
	}
	return messages, nil
}

func (c *ChannelMessage) GetMentionedUsernames() []string {
	flattened := make([]string, 0)

	res := mentionRegex.FindAllStringSubmatch(c.Body, -1)
	if len(res) == 0 {
		return flattened
	}

	participants := map[string]struct{}{}
	// remove duplicate mentions
	for _, ele := range res {
		participants[ele[1]] = struct{}{}
	}

	for participant := range participants {
		flattened = append(flattened, participant)
	}

	return flattened
}

func (c *ChannelMessage) FetchTotalMessageCount(q *Query) (int, error) {
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id":    q.AccountId,
			"type_constant": q.Type,
		},
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
	}

	return c.CountWithQuery(query)
}

func (c *ChannelMessage) FetchMessageIds(q *Query) ([]int64, error) {
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id":    q.AccountId,
			"type_constant": q.Type,
		},
		Pluck:      "id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort: map[string]string{
			"created_at": "DESC",
		},
	}

	var messageIds []int64
	if err := c.Some(&messageIds, query); err != nil {
		return nil, err
	}

	if messageIds == nil {
		return make([]int64, 0), nil
	}

	return messageIds, nil
}

// BySlug fetchs channel message by its slug
// checks if message is in the channel or not
func (c *ChannelMessage) BySlug(query *Query) error {
	if query.Slug == "" {
		return errors.New("slug is not set")
	}

	// fetch message itself
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"slug": query.Slug,
		},
	}

	if err := c.One(q); err != nil {
		return nil
	}

	// fetch channel by group name
	query.Name = query.GroupName
	query.Type = Channel_TYPE_GROUP
	ch := NewChannel()
	channel, err := ch.ByName(query)
	if err != nil {
		return err
	}

	if channel.Id == 0 {
		return errors.New("channel is not found")
	}

	// check if message is in the channel
	cml := NewChannelMessageList()
	res, err := cml.IsInChannel(c.Id, channel.Id)
	if err != nil {
		return err
	}

	// if message is not in the channel
	if !res {
		return gorm.RecordNotFound
	}

	return nil
}
<|endoftext|>


<commit_before>package model

import "testing"

func TestUserChoiceAnswer(t *testing.T) {
	choicesDisplay := []string{"エラーが出てる", "何もしてないけど壊れた", "画面が映らない", "インターネットに繋がらない"}
	choices := []string{
		"エラーが出てる,えらーがでてる",
		"何もしてないけど壊れた,なにもしてないけどこわれた",
		"画面が映らない,がめんがうつらない",
		"インターネットに繋がらない,いんたーねっとにつならがない",
	}
	userAnswer := "何もしてないけど壊れたよ"
	topRateIndex := UserChoiceAnswer(choices, userAnswer)
	if topRateIndex != 1 {
		t.Fatalf("%sはchoicesの1番目が選択されるべきです 実際の値は%d", userAnswer, topRateIndex)
	}
	if choicesDisplay[topRateIndex] != choicesDisplay[1] {
		t.Fatalf("選び出したい選択肢は%sです 実際の値は%s", choicesDisplay[1], choicesDisplay[topRateIndex])
	}

	userAnswer2 := "ネットに繋がらない"
	topRateIndex2 := UserChoiceAnswer(choices, userAnswer2)
	if topRateIndex2 != 3 {
		t.Fatalf("%sはchoicesの3番目が選択されるべきです 実際の値は%d", userAnswer, topRateIndex)
	}
	if choicesDisplay[topRateIndex2] != choicesDisplay[3] {
		t.Fatalf("選び出したい選択肢は%sです 実際の値は%s", choicesDisplay[3], choicesDisplay[topRateIndex])
	}
}
<commit_msg>fix: lint<commit_after>package model

import "testing"

func TestUserChoiceAnswer(t *testing.T) {
	choicesDisplay := []string{"エラーが出てる", "何もしてないけど壊れた", "画面が映らない", "インターネットに繋がらない"}
	choices := []string{
		"エラーが出てる,えらーがでてる",
		"何もしてないけど壊れた,なにもしてないけどこわれた",
		"画面が映らない,がめんがうつらない",
		"インターネットに繋がらない,いんたーねっとにつならがない",
	}
	userAnswer := "何もしてないけど壊れたよ"
	topRateIndex, err := UserChoiceAnswer(choices, userAnswer)
	if err != nil {
		t.Fatalf("エラーが発生しました %v", err)
	}
	if topRateIndex != 1 {
		t.Fatalf("%sはchoicesの1番目が選択されるべきです 実際の値は%d", userAnswer, topRateIndex)
	}
	if choicesDisplay[topRateIndex] != choicesDisplay[1] {
		t.Fatalf("選び出したい選択肢は%sです 実際の値は%s", choicesDisplay[1], choicesDisplay[topRateIndex])
	}

	userAnswer2 := "ネットに繋がらない"
	topRateIndex2, err := UserChoiceAnswer(choices, userAnswer2)
	if err != nil {
		t.Fatalf("エラーが発生しました %v", err)
	}
	if topRateIndex2 != 3 {
		t.Fatalf("%sはchoicesの3番目が選択されるべきです 実際の値は%d", userAnswer, topRateIndex)
	}
	if choicesDisplay[topRateIndex2] != choicesDisplay[3] {
		t.Fatalf("選び出したい選択肢は%sです 実際の値は%s", choicesDisplay[3], choicesDisplay[topRateIndex])
	}
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com
package main

import (
	"bufio"
	"encoding/xml"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"oval"
	"path"
	"regexp"
	"strings"
)

type cveEntry struct {
	cveID         string
	pkgName       string
	patchVersions map[string]string
}

var entries []cveEntry
var matchFilter *regexp.Regexp

func parseEntryFile(fpath string) (ret cveEntry) {
	const (
		_ = iota
		INNER_NONE
		INNER_PATCH
	)
	fd, err := os.Open(fpath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	defer func() {
		fd.Close()
	}()

	scanner := bufio.NewScanner(fd)
	parserMode := INNER_NONE
	curPkgName := ""
	ret.cveID = path.Base(fpath)
	ret.patchVersions = make(map[string]string)
	for scanner.Scan() {
		tokens := strings.Fields(scanner.Text())
		if len(tokens) == 0 {
			parserMode = INNER_NONE
			curPkgName = ""
			continue
		}

		if strings.HasPrefix(tokens[0], "Patches_") {
			parserMode = INNER_PATCH
			curPkgName = strings.TrimPrefix(tokens[0], "Patches_")
			curPkgName = strings.TrimRight(curPkgName, ":")
			ret.pkgName = curPkgName
			continue
		}

		if parserMode == INNER_PATCH {
			if len(tokens) < 2 {
				continue
			}
			idx := strings.Index(tokens[0], "_")
			if idx == -1 {
				continue
			}
			distname := tokens[0][:idx]
			if tokens[1] == "released" && len(tokens) > 2 {
				patchver := tokens[2]
				patchver = strings.Trim(patchver, "()")
				ret.patchVersions[distname] = patchver
			}
		}
	}
	if err = scanner.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	return
}

func loadEntries(dirpath string) {
	dirents, err := ioutil.ReadDir(dirpath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	for _, i := range dirents {
		if !strings.HasPrefix(i.Name(), "CVE-") {
			continue
		}
		if matchFilter != nil {
			if !matchFilter.MatchString(i.Name()) {
				continue
			}
		}
		fname := path.Join(dirpath, i.Name())
		entries = append(entries, parseEntryFile(fname))
	}
}

func addDefinition(o *oval.GOvalDefinitions, prefix string, dist string, cve cveEntry) {
	// Create a state
	stateid := fmt.Sprintf("%v-state", prefix)
	state := oval.GDPKGInfoState{}
	state.ID = stateid
	state.EVRCheck.DataType = "evr_string"
	state.EVRCheck.Operation = "less than"
	state.EVRCheck.Value = cve.patchVersions[dist]

	// Create an object definition for the package
	objid := fmt.Sprintf("%v-object", prefix)
	obj := oval.GDPKGInfoObj{}
	obj.Name = cve.pkgName
	obj.ID = objid

	o.States.DPKGInfoStates = append(o.States.DPKGInfoStates, state)
	o.Objects.DPKGInfoObjects = append(o.Objects.DPKGInfoObjects, obj)
}

func processEntries() {
	root := oval.GOvalDefinitions{}

	for i, ent := range entries {
		prefix := fmt.Sprintf("ubuntu-%v", i)
		for x := range ent.patchVersions {
			addDefinition(&root, prefix, x, ent)
		}
	}

	enc := xml.NewEncoder(os.Stdout)
	enc.Indent("", "  ")
	if err := enc.Encode(root); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}

func main() {
	var fMatch string

	flag.StringVar(&fMatch, "i", "", "filter regexp")
	flag.Parse()
	args := flag.Args()
	if len(args) < 1 {
		fmt.Fprintf(os.Stderr, "specify path to ubuntu-cve-tracker directory\n")
		os.Exit(1)
	}

	fm, err := regexp.Compile(fMatch)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	matchFilter = fm

	entries = make([]cveEntry, 0)

	procdir := path.Join(args[0], "active")
	loadEntries(procdir)
	procdir = path.Join(args[0], "retired")
	loadEntries(procdir)
	processEntries()
}
<commit_msg>link package object and state in test<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com
package main

import (
	"bufio"
	"encoding/xml"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"oval"
	"path"
	"regexp"
	"strings"
)

type cveEntry struct {
	cveID         string
	pkgName       string
	patchVersions map[string]string
}

var entries []cveEntry
var matchFilter *regexp.Regexp

func parseEntryFile(fpath string) (ret cveEntry) {
	const (
		_ = iota
		INNER_NONE
		INNER_PATCH
	)
	fd, err := os.Open(fpath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	defer func() {
		fd.Close()
	}()

	scanner := bufio.NewScanner(fd)
	parserMode := INNER_NONE
	curPkgName := ""
	ret.cveID = path.Base(fpath)
	ret.patchVersions = make(map[string]string)
	for scanner.Scan() {
		tokens := strings.Fields(scanner.Text())
		if len(tokens) == 0 {
			parserMode = INNER_NONE
			curPkgName = ""
			continue
		}

		if strings.HasPrefix(tokens[0], "Patches_") {
			parserMode = INNER_PATCH
			curPkgName = strings.TrimPrefix(tokens[0], "Patches_")
			curPkgName = strings.TrimRight(curPkgName, ":")
			ret.pkgName = curPkgName
			continue
		}

		if parserMode == INNER_PATCH {
			if len(tokens) < 2 {
				continue
			}
			idx := strings.Index(tokens[0], "_")
			if idx == -1 {
				continue
			}
			distname := tokens[0][:idx]
			if tokens[1] == "released" && len(tokens) > 2 {
				patchver := tokens[2]
				patchver = strings.Trim(patchver, "()")
				ret.patchVersions[distname] = patchver
			}
		}
	}
	if err = scanner.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	return
}

func loadEntries(dirpath string) {
	dirents, err := ioutil.ReadDir(dirpath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	for _, i := range dirents {
		if !strings.HasPrefix(i.Name(), "CVE-") {
			continue
		}
		if matchFilter != nil {
			if !matchFilter.MatchString(i.Name()) {
				continue
			}
		}
		fname := path.Join(dirpath, i.Name())
		entries = append(entries, parseEntryFile(fname))
	}
}

func addDefinition(o *oval.GOvalDefinitions, prefix string, dist string, cve cveEntry) {
	// Create a state
	stateid := fmt.Sprintf("%v-state", prefix)
	state := oval.GDPKGInfoState{}
	state.ID = stateid
	state.EVRCheck.DataType = "evr_string"
	state.EVRCheck.Operation = "less than"
	state.EVRCheck.Value = cve.patchVersions[dist]

	// Create an object definition for the package
	objid := fmt.Sprintf("%v-object", prefix)
	obj := oval.GDPKGInfoObj{}
	obj.Name = cve.pkgName
	obj.ID = objid

	// Create a test
	testid := fmt.Sprintf("%v-test", prefix)
	test := oval.GDPKGInfoTest{}
	test.ID = testid
	test.Object.ObjectRef = objid
	test.State.StateRef = stateid

	o.States.DPKGInfoStates = append(o.States.DPKGInfoStates, state)
	o.Objects.DPKGInfoObjects = append(o.Objects.DPKGInfoObjects, obj)
	o.Tests.DPKGInfoTests = append(o.Tests.DPKGInfoTests, test)
}

func processEntries() {
	root := oval.GOvalDefinitions{}

	for i, ent := range entries {
		prefix := fmt.Sprintf("ubuntu-%v", i)
		for x := range ent.patchVersions {
			addDefinition(&root, prefix, x, ent)
		}
	}

	enc := xml.NewEncoder(os.Stdout)
	enc.Indent("", "  ")
	if err := enc.Encode(root); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}

func main() {
	var fMatch string

	flag.StringVar(&fMatch, "i", "", "filter regexp")
	flag.Parse()
	args := flag.Args()
	if len(args) < 1 {
		fmt.Fprintf(os.Stderr, "specify path to ubuntu-cve-tracker directory\n")
		os.Exit(1)
	}

	fm, err := regexp.Compile(fMatch)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
	matchFilter = fm

	entries = make([]cveEntry, 0)

	procdir := path.Join(args[0], "active")
	loadEntries(procdir)
	procdir = path.Join(args[0], "retired")
	loadEntries(procdir)
	processEntries()
}
<|endoftext|>


<commit_before>package data

import (
	"log"
	"testing"

	"github.com/mdlayher/goat/goat/common"
)

// TestWhitelistRecord verifies that WhitelistRecord creation, methods, save, load, and delete work properly
func TestWhitelistRecord(t *testing.T) {
	log.Println("TestWhitelistRecord()")

	// Load config
	config, err := common.LoadConfig()
	if err != nil {
		t.Fatalf("Could not load configuration: %s", err.Error())
	}
	common.Static.Config = config

	// Generate mock WhitelistRecord
	whitelist := WhitelistRecord{
		Client: "goat_test",
		Approved: true,
	}

	// Save mock whitelist
	if err := whitelist.Save(); err != nil {
		t.Fatalf("Failed to save mock whitelist: %s", err.Error())
	}

	// Load mock whitelist to fetch ID
	whitelist, err = whitelist.Load(whitelist.Client, "client")
	if whitelist == (WhitelistRecord{}) || err != nil {
		t.Fatalf("Failed to load mock whitelist: %s", err.Error())
	}

	// Delete mock whitelist
	if err := whitelist.Delete(); err != nil {
		t.Fatalf("Failed to delete mock whitelist: %s", err.Error())
	}
}
<commit_msg>go fmt<commit_after>package data

import (
	"log"
	"testing"

	"github.com/mdlayher/goat/goat/common"
)

// TestWhitelistRecord verifies that WhitelistRecord creation, methods, save, load, and delete work properly
func TestWhitelistRecord(t *testing.T) {
	log.Println("TestWhitelistRecord()")

	// Load config
	config, err := common.LoadConfig()
	if err != nil {
		t.Fatalf("Could not load configuration: %s", err.Error())
	}
	common.Static.Config = config

	// Generate mock WhitelistRecord
	whitelist := WhitelistRecord{
		Client:   "goat_test",
		Approved: true,
	}

	// Save mock whitelist
	if err := whitelist.Save(); err != nil {
		t.Fatalf("Failed to save mock whitelist: %s", err.Error())
	}

	// Load mock whitelist to fetch ID
	whitelist, err = whitelist.Load(whitelist.Client, "client")
	if whitelist == (WhitelistRecord{}) || err != nil {
		t.Fatalf("Failed to load mock whitelist: %s", err.Error())
	}

	// Delete mock whitelist
	if err := whitelist.Delete(); err != nil {
		t.Fatalf("Failed to delete mock whitelist: %s", err.Error())
	}
}
<|endoftext|>


<commit_before>package tcpcall

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"math/rand"
	"testing"
	"time"
)

func TestClientConcurrency(t *testing.T) {
	// configure
	concurrency := 500
	// create server
	serverConf := NewServerConf()
	serverConf.PortNumber = 6000
	serverConf.MaxConnections = 1
	serverConf.Concurrency = concurrency
	serverConf.RequestCallback = func(data []byte) []byte {
		return serve(t, data)
	}
	server, err := Listen(serverConf)
	if err != nil {
		t.Fatalf("listen: %s", err)
	}
	// create client
	clientConf := NewClientConf()
	clientConf.Concurrency = concurrency
	client, err := Dial("127.0.0.1:6000", clientConf)
	if err != nil {
		t.Fatalf("dial: %s", err)
	}
	// issue a lot of requests in parallel
	chn := make(chan int, concurrency)
	for i := 0; i < concurrency; i++ { // spawn senders
		go workerLoop(t, chn, client)
	}
	for i := 0; i < 500000; i++ { // send communicate signals
		chn <- i
	}
	for i := 0; i < concurrency; i++ { // send termination signals
		chn <- -1
	}
	for 0 < len(chn) { // wait for senders to terminate
		time.Sleep(10 * time.Millisecond)
	}
	// cleanup
	client.Close()
	server.Stop()
}

func workerLoop(t *testing.T, chn chan int, client *Client) {
	for sig := range chn {
		if sig < 0 {
			return
		}
		req := genReq(sig)
		rep, err := client.Req(req, time.Second)
		if err != nil {
			t.Fatalf("req #%d failed: %s", sig, err)
		}
		if !bytes.Equal(req, rep) {
			t.Fatalf("req #%d: assert failed: %v != %v", sig, req, rep)
		}
	}
}

func serve(t *testing.T, data []byte) []byte {
	if err := decodeReq(data); err != nil {
		t.Errorf("server: failed to decode %v: %v", data, err)
	}
	return data
}

func genReq(id int) []byte {
	size := rand.Intn(256)
	msg := make([]byte, 5+size)
	binary.BigEndian.PutUint32(msg, uint32(id))
	msg[4] = byte(size)
	for i := 0; i < size; i++ {
		msg[5+i] = byte(size + id)
	}
	return msg
}

func decodeReq(data []byte) error {
	if len(data) < 5 {
		return errors.New("too short header")
	}
	id := int(binary.BigEndian.Uint32(data))
	size := int(data[4])
	if len(data) != 5+size {
		return fmt.Errorf("size mismatch. Expected %d but found %d",
			5+size, len(data))
	}
	expect := byte(size + id)
	for i := 0; i < size; i++ {
		if data[5+i] != expect {
			return fmt.Errorf("byte #%d. Expected %d but found %d",
				5+i, expect, data[5+i])
		}
	}
	return nil
}
<commit_msg>Refactor client concurrency test<commit_after>package tcpcall

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"testing"
	"time"
)

func TestClientConcurrency(t *testing.T) {
	// configure
	concurrency := 500
	// create server
	serverConf := NewServerConf()
	serverConf.PortNumber = 6000
	serverConf.MaxConnections = 1
	serverConf.Concurrency = concurrency
	serverConf.RequestCallback = func(data []byte) []byte {
		return serve(t, data)
	}
	server, err := Listen(serverConf)
	if err != nil {
		t.Fatalf("listen: %s", err)
	}
	// create client
	clientConf := NewClientConf()
	clientConf.Concurrency = concurrency
	client, err := Dial("127.0.0.1:6000", clientConf)
	if err != nil {
		t.Fatalf("dial: %s", err)
	}
	// issue a lot of requests in parallel
	var (
		wg  sync.WaitGroup
		chn = make(chan int, concurrency)
	)
	log.Printf("spawning...")
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go workerLoop(t, &wg, chn, client)
	}
	log.Printf("communicating...")
	for i := 0; i < 500000; i++ {
		chn <- i
	}
	log.Printf("terminating...")
	close(chn)
	wg.Wait()
	log.Printf("cleanup...")
	client.Close()
	server.Stop()
}

func workerLoop(t *testing.T, wg *sync.WaitGroup, chn chan int, client *Client) {
	defer wg.Done()
	for {
		sig, ok := <-chn
		if !ok {
			return
		}
		req := genReq(sig)
		rep, err := client.Req(req, 2*time.Second)
		if err != nil {
			t.Fatalf("req #%d failed: %s", sig, err)
		}
		if !bytes.Equal(req, rep) {
			t.Fatalf("req #%d: assert failed: %v != %v", sig, req, rep)
		}
	}
}

func serve(t *testing.T, data []byte) []byte {
	if err := decodeReq(data); err != nil {
		t.Errorf("server: failed to decode %v: %v", data, err)
	}
	return data
}

func genReq(id int) []byte {
	size := rand.Intn(256)
	msg := make([]byte, 5+size)
	binary.BigEndian.PutUint32(msg, uint32(id))
	msg[4] = byte(size)
	for i := 0; i < size; i++ {
		msg[5+i] = byte(size + id)
	}
	return msg
}

func decodeReq(data []byte) error {
	if len(data) < 5 {
		return errors.New("too short header")
	}
	id := int(binary.BigEndian.Uint32(data))
	size := int(data[4])
	if len(data) != 5+size {
		return fmt.Errorf("size mismatch. Expected %d but found %d",
			5+size, len(data))
	}
	expect := byte(size + id)
	for i := 0; i < size; i++ {
		if data[5+i] != expect {
			return fmt.Errorf("byte #%d. Expected %d but found %d",
				5+i, expect, data[5+i])
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package tcpcall

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"testing"
	"time"
)

func TestClientConcurrency(t *testing.T) {
	// configure
	concurrency := 500
	// create server
	serverConf := NewServerConf()
	serverConf.PortNumber = 6000
	serverConf.MaxConnections = 1
	serverConf.Concurrency = concurrency
	serverConf.RequestCallback = func(data []byte) []byte {
		return serve(t, data)
	}
	server, err := Listen(serverConf)
	if err != nil {
		t.Fatalf("listen: %s", err)
	}
	// create client
	clientConf := NewClientConf()
	clientConf.Concurrency = concurrency
	client, err := Dial("127.0.0.1:6000", clientConf)
	if err != nil {
		t.Fatalf("dial: %s", err)
	}
	// issue a lot of requests in parallel
	var (
		wg  sync.WaitGroup
		chn = make(chan int, concurrency)
	)
	log.Printf("spawning...")
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go workerLoop(t, &wg, chn, client)
	}
	log.Printf("communicating...")
	for i := 0; i < 500000; i++ {
		chn <- i
	}
	log.Printf("terminating...")
	close(chn)
	wg.Wait()
	log.Printf("cleanup...")
	client.Close()
	server.Stop()
}

func workerLoop(t *testing.T, wg *sync.WaitGroup, chn chan int, client *Client) {
	defer wg.Done()
	for {
		sig, ok := <-chn
		if !ok {
			return
		}
		req := genReq(sig)
		rep, err := client.Req(req, 2*time.Second)
		if err != nil {
			t.Fatalf("req #%d failed: %s", sig, err)
		}
		if !bytes.Equal(req, rep) {
			t.Fatalf("req #%d: assert failed: %v != %v", sig, req, rep)
		}
	}
}

func serve(t *testing.T, data []byte) []byte {
	if err := decodeReq(data); err != nil {
		t.Errorf("server: failed to decode %v: %v", data, err)
	}
	return data
}

func genReq(id int) []byte {
	size := rand.Intn(256)
	msg := make([]byte, 5+size)
	binary.BigEndian.PutUint32(msg, uint32(id))
	msg[4] = byte(size)
	for i := 0; i < size; i++ {
		msg[5+i] = byte(size + id)
	}
	return msg
}

func decodeReq(data []byte) error {
	if len(data) < 5 {
		return errors.New("too short header")
	}
	id := int(binary.BigEndian.Uint32(data))
	size := int(data[4])
	if len(data) != 5+size {
		return fmt.Errorf("size mismatch. Expected %d but found %d",
			5+size, len(data))
	}
	expect := byte(size + id)
	for i := 0; i < size; i++ {
		if data[5+i] != expect {
			return fmt.Errorf("byte #%d. Expected %d but found %d",
				5+i, expect, data[5+i])
		}
	}
	return nil
}
<commit_msg>Minor test optimization<commit_after>package tcpcall

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"testing"
	"time"
)

func TestClientConcurrency(t *testing.T) {
	// configure
	concurrency := 500
	// create server
	serverConf := NewServerConf()
	serverConf.PortNumber = 6000
	serverConf.MaxConnections = 1
	serverConf.Concurrency = concurrency
	serverConf.RequestCallback = func(data []byte) []byte {
		return serve(t, data)
	}
	server, err := Listen(serverConf)
	if err != nil {
		t.Fatalf("listen: %s", err)
	}
	// create client
	clientConf := NewClientConf()
	clientConf.Concurrency = concurrency
	client, err := Dial("127.0.0.1:6000", clientConf)
	if err != nil {
		t.Fatalf("dial: %s", err)
	}
	// issue a lot of requests in parallel
	var (
		wg  sync.WaitGroup
		chn = make(chan int, concurrency)
	)
	log.Printf("spawning...")
	wg.Add(concurrency)
	for i := 0; i < concurrency; i++ {
		go workerLoop(t, &wg, chn, client)
	}
	log.Printf("communicating...")
	for i := 0; i < 500000; i++ {
		chn <- i
	}
	log.Printf("terminating...")
	close(chn)
	wg.Wait()
	log.Printf("cleanup...")
	client.Close()
	server.Stop()
}

func workerLoop(t *testing.T, wg *sync.WaitGroup, chn chan int, client *Client) {
	defer wg.Done()
	for {
		sig, ok := <-chn
		if !ok {
			return
		}
		req := genReq(sig)
		rep, err := client.Req(req, 2*time.Second)
		if err != nil {
			t.Fatalf("req #%d failed: %s", sig, err)
		}
		if !bytes.Equal(req, rep) {
			t.Fatalf("req #%d: assert failed: %v != %v", sig, req, rep)
		}
	}
}

func serve(t *testing.T, data []byte) []byte {
	if err := decodeReq(data); err != nil {
		t.Errorf("server: failed to decode %v: %v", data, err)
	}
	return data
}

func genReq(id int) []byte {
	size := rand.Intn(256)
	msg := make([]byte, 5+size)
	binary.BigEndian.PutUint32(msg, uint32(id))
	msg[4] = byte(size)
	for i := 0; i < size; i++ {
		msg[5+i] = byte(size + id)
	}
	return msg
}

func decodeReq(data []byte) error {
	if len(data) < 5 {
		return errors.New("too short header")
	}
	id := int(binary.BigEndian.Uint32(data))
	size := int(data[4])
	if len(data) != 5+size {
		return fmt.Errorf("size mismatch. Expected %d but found %d",
			5+size, len(data))
	}
	expect := byte(size + id)
	for i := 0; i < size; i++ {
		if data[5+i] != expect {
			return fmt.Errorf("byte #%d. Expected %d but found %d",
				5+i, expect, data[5+i])
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package goppetto

import (
	"testing"
)

var ed EventDispatcher
var isCalled bool
var done = make(chan bool)

var f = func(emsg *EventMessage) *EventMessage {
	isCalled = true
	done <- true

	return emsg
}

func setUp() {
	ed = EventDispatcher{
		callbacks: make(map[string][]func(*EventMessage) *EventMessage),
	}
	isCalled = false
}

func TestBind(t *testing.T) {
	setUp()
	ed.Bind("some_event", f)
	ed.Bind("some_event", f)

	if len(ed.callbacks["some_event"]) != 2 {
		t.Errorf("EventDispatcher should have 2 callbacks bound to 'some_event' got %v", len(ed.callbacks["some_event"]))
	}
}

func TestListen(t *testing.T) {
	setUp()
	msg := `{"event": "some_event", "data": {"pin_id": 1, "state": 0}}`
	messages := make(chan string, 1)

	ed.Bind("some_event", f)
	go ed.Listen(messages)
	messages <- msg

	<-done

	if isCalled == false {
		t.Error("someEvent has not been dispatched.")
	}
}

func TestDispatch(t *testing.T) {
	setUp()
	msg := EventMessage{Event: "some_event"}

	ed.Bind("some_event", f)
	ed.Dispatch(&msg)

	// Wait for the someEvent to be called.
	<-done

	if isCalled == false {
		t.Error("Method has not been dispatched.")
	}
}
<commit_msg>#10 Rewrote tests with Go Convey.<commit_after>package goppetto

import (
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func TestEventDispatcher(t *testing.T) {
	Convey("Given I have a EventDispatcher", t, func() {
		i := 0
		done := make(chan bool)

		em := EventMessage{"pin_state", make(map[string]interface{})}

		ed := EventDispatcher{make(map[string][]func(*EventMessage) *EventMessage)}

		Convey("When I bind a callback to an event", nil)

		Convey("And that the EventDispatcher receives this event", func() {
			ed.Bind("pin_state", func(e *EventMessage) *EventMessage {
				i += 1
				done <- true

				return e
			})

			ed.Dispatch(&em)

			Convey("Then the callback must be executed.", func() {
				<-done
				So(i, ShouldEqual, 1)
			})
		})

		Convey("When I bind multiple callbacks to an event", nil)

		Convey("And the EventDispatcher receives this event", func() {
			signal := make(chan bool)

			ed.Bind("pin_state", func(e *EventMessage) *EventMessage {
				i += 1
				signal <- true
				return e
			})

			// This callback receives a value in channel `signal`. This value
			// can only be send by the previous callback. If both callbacks
			// where executed seqentually this construct would cause a dead
			// lock.
			ed.Bind("pin_state", func(e *EventMessage) *EventMessage {
				i += 1
				<-signal
				done <- true
				return e
			})

			ed.Dispatch(&em)

			Convey("Then the callbacks must be executed in parallel.", func() {
				<-done
				So(i, ShouldEqual, 2)
			})
		})
	})
}
<|endoftext|>


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

package sqlparser

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/wandoulabs/cm/hack"
	"github.com/wandoulabs/cm/sqltypes"
)

const EOFCHAR = 0x100

// Tokenizer is the struct used to generate SQL
// tokens for the parser.
type Tokenizer struct {
	InStream      *strings.Reader
	AllowComments bool
	ForceEOF      bool
	lastChar      uint16
	Position      int
	errorToken    []byte
	LastError     string
	posVarIndex   int
	ParseTree     Statement
}

// NewStringTokenizer creates a new Tokenizer for the
// sql string.
func NewStringTokenizer(sql string) *Tokenizer {
	return &Tokenizer{InStream: strings.NewReader(sql)}
}

var keywords = map[string]int{
	"all":           ALL,
	"alter":         ALTER,
	"analyze":       ANALYZE,
	"and":           AND,
	"as":            AS,
	"asc":           ASC,
	"between":       BETWEEN,
	"by":            BY,
	"case":          CASE,
	"create":        CREATE,
	"cross":         CROSS,
	"default":       DEFAULT,
	"delete":        DELETE,
	"desc":          DESC,
	"describe":      DESCRIBE,
	"distinct":      DISTINCT,
	"drop":          DROP,
	"duplicate":     DUPLICATE,
	"else":          ELSE,
	"end":           END,
	"except":        EXCEPT,
	"exists":        EXISTS,
	"explain":       EXPLAIN,
	"for":           FOR,
	"force":         FORCE,
	"from":          FROM,
	"group":         GROUP,
	"having":        HAVING,
	"if":            IF,
	"ignore":        IGNORE,
	"in":            IN,
	"index":         INDEX,
	"inner":         INNER,
	"insert":        INSERT,
	"intersect":     INTERSECT,
	"into":          INTO,
	"is":            IS,
	"join":          JOIN,
	"key":           KEY,
	"left":          LEFT,
	"like":          LIKE,
	"limit":         LIMIT,
	"lock":          LOCK,
	"minus":         MINUS,
	"natural":       NATURAL,
	"not":           NOT,
	"null":          NULL,
	"on":            ON,
	"or":            OR,
	"order":         ORDER,
	"outer":         OUTER,
	"rename":        RENAME,
	"right":         RIGHT,
	"select":        SELECT,
	"set":           SET,
	"show":          SHOW,
	"straight_join": STRAIGHT_JOIN,
	"table":         TABLE,
	"then":          THEN,
	"to":            TO,
	"union":         UNION,
	"unique":        UNIQUE,
	"update":        UPDATE,
	"use":           USE,
	"using":         USING,
	"values":        VALUES,
	"view":          VIEW,
	"when":          WHEN,
	"where":         WHERE,
}

// Lex returns the next token form the Tokenizer.
// This function is used by go yacc.
func (tkn *Tokenizer) Lex(lval *yySymType) int {
	typ, val := tkn.Scan()
	for typ == COMMENT {
		if tkn.AllowComments {
			break
		}
		typ, val = tkn.Scan()
	}
	switch typ {
	case ID, STRING, NUMBER, VALUE_ARG, LIST_ARG, COMMENT:
		lval.bytes = val
	}
	tkn.errorToken = val
	return typ
}

// Error is called by go yacc if there's a parsing error.
func (tkn *Tokenizer) Error(err string) {
	buf := bytes.NewBuffer(make([]byte, 0, 32))
	if tkn.errorToken != nil {
		fmt.Fprintf(buf, "%s at position %v near %s", err, tkn.Position, tkn.errorToken)
	} else {
		fmt.Fprintf(buf, "%s at position %v", err, tkn.Position)
	}
	tkn.LastError = buf.String()
}

func (tkn *Tokenizer) scanHexValue() (int, []byte) {
	return 0, nil
}

// Scan scans the tokenizer for the next token and returns
// the token type and an optional value.
func (tkn *Tokenizer) Scan() (int, []byte) {
	if tkn.ForceEOF {
		return 0, nil
	}

	if tkn.lastChar == 0 {
		tkn.next()
	}
	tkn.skipBlank()
	switch ch := tkn.lastChar; {
	case isLetter(ch):
		if ch == 'x' {
			tkn.next()
			c := tkn.lastChar
			if c == '\'' {
				return tkn.scanHexValue()
			}
		}
		return tkn.scanIdentifier()
	case isDigit(ch):
		return tkn.scanNumber(false)
	case ch == ':':
		return tkn.scanBindVar()
	default:
		tkn.next()
		switch ch {
		case EOFCHAR:
			return 0, nil
		case '=', ',', ';', '(', ')', '+', '*', '%', '&', '|', '^', '~':
			return int(ch), nil
		case '?':
			tkn.posVarIndex++
			buf := new(bytes.Buffer)
			fmt.Fprintf(buf, ":v%d", tkn.posVarIndex)
			return VALUE_ARG, buf.Bytes()
		case '.':
			if isDigit(tkn.lastChar) {
				return tkn.scanNumber(true)
			} else {
				return int(ch), nil
			}
		case '/':
			switch tkn.lastChar {
			case '/':
				tkn.next()
				return tkn.scanCommentType1("//")
			case '*':
				tkn.next()
				return tkn.scanCommentType2()
			default:
				return int(ch), nil
			}
		case '-':
			if tkn.lastChar == '-' {
				tkn.next()
				return tkn.scanCommentType1("--")
			} else {
				return int(ch), nil
			}
		case '<':
			switch tkn.lastChar {
			case '>':
				tkn.next()
				return NE, nil
			case '=':
				tkn.next()
				switch tkn.lastChar {
				case '>':
					tkn.next()
					return NULL_SAFE_EQUAL, nil
				default:
					return LE, nil
				}
			default:
				return int(ch), nil
			}
		case '>':
			if tkn.lastChar == '=' {
				tkn.next()
				return GE, nil
			} else {
				return int(ch), nil
			}
		case '!':
			if tkn.lastChar == '=' {
				tkn.next()
				return NE, nil
			} else {
				return LEX_ERROR, []byte("!")
			}
		case '\'', '"':
			return tkn.scanString(ch, STRING)
		case '`':
			return tkn.scanLiteralIdentifier()
		default:
			return LEX_ERROR, []byte{byte(ch)}
		}
	}
}

func (tkn *Tokenizer) skipBlank() {
	ch := tkn.lastChar
	for ch == ' ' || ch == '\n' || ch == '\r' || ch == '\t' {
		tkn.next()
		ch = tkn.lastChar
	}
}

func (tkn *Tokenizer) scanIdentifier() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	buffer.WriteByte(byte(tkn.lastChar))
	for tkn.next(); isLetter(tkn.lastChar) || isDigit(tkn.lastChar); tkn.next() {
		buffer.WriteByte(byte(tkn.lastChar))
	}
	lowered := bytes.ToLower(buffer.Bytes())
	if keywordId, found := keywords[hack.String(lowered)]; found {
		return keywordId, lowered
	}

	return ID, buffer.Bytes()
}

func (tkn *Tokenizer) scanLiteralIdentifier() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteByte(byte(tkn.lastChar))
	if !isLetter(tkn.lastChar) {
		return LEX_ERROR, buffer.Bytes()
	}
	for tkn.next(); isLetter(tkn.lastChar) || isDigit(tkn.lastChar); tkn.next() {
		buffer.WriteByte(byte(tkn.lastChar))
	}
	if tkn.lastChar != '`' {
		return LEX_ERROR, buffer.Bytes()
	}
	tkn.next()
	return ID, buffer.Bytes()
}

func (tkn *Tokenizer) scanBindVar() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteByte(byte(tkn.lastChar))
	token := VALUE_ARG
	tkn.next()
	if tkn.lastChar == ':' {
		token = LIST_ARG
		buffer.WriteByte(byte(tkn.lastChar))
		tkn.next()
	}
	if !isLetter(tkn.lastChar) {
		return LEX_ERROR, buffer.Bytes()
	}
	for isLetter(tkn.lastChar) || isDigit(tkn.lastChar) || tkn.lastChar == '.' {
		buffer.WriteByte(byte(tkn.lastChar))
		tkn.next()
	}
	return token, buffer.Bytes()
}

func (tkn *Tokenizer) scanMantissa(base int, buffer *bytes.Buffer) {
	for digitVal(tkn.lastChar) < base {
		tkn.ConsumeNext(buffer)
	}
}

func (tkn *Tokenizer) scanNumber(seenDecimalPoint bool) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	if seenDecimalPoint {
		buffer.WriteByte('.')
		tkn.scanMantissa(10, buffer)
		goto exponent
	}

	if tkn.lastChar == '0' {
		// int or float
		tkn.ConsumeNext(buffer)
		if tkn.lastChar == 'x' || tkn.lastChar == 'X' {
			// hexadecimal int
			tkn.ConsumeNext(buffer)
			tkn.scanMantissa(16, buffer)
		} else {
			// octal int or float
			seenDecimalDigit := false
			tkn.scanMantissa(8, buffer)
			if tkn.lastChar == '8' || tkn.lastChar == '9' {
				// illegal octal int or float
				seenDecimalDigit = true
				tkn.scanMantissa(10, buffer)
			}
			if tkn.lastChar == '.' || tkn.lastChar == 'e' || tkn.lastChar == 'E' {
				goto fraction
			}
			// octal int
			if seenDecimalDigit {
				return LEX_ERROR, buffer.Bytes()
			}
		}
		goto exit
	}

	// decimal int or float
	tkn.scanMantissa(10, buffer)

fraction:
	if tkn.lastChar == '.' {
		tkn.ConsumeNext(buffer)
		tkn.scanMantissa(10, buffer)
	}

exponent:
	if tkn.lastChar == 'e' || tkn.lastChar == 'E' {
		tkn.ConsumeNext(buffer)
		if tkn.lastChar == '+' || tkn.lastChar == '-' {
			tkn.ConsumeNext(buffer)
		}
		tkn.scanMantissa(10, buffer)
	}

exit:
	return NUMBER, buffer.Bytes()
}

func (tkn *Tokenizer) scanString(delim uint16, typ int) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	for {
		ch := tkn.lastChar
		tkn.next()
		if ch == delim {
			if tkn.lastChar == delim {
				tkn.next()
			} else {
				break
			}
		} else if ch == '\\' {
			if tkn.lastChar == EOFCHAR {
				return LEX_ERROR, buffer.Bytes()
			}
			if decodedChar := sqltypes.SqlDecodeMap[byte(tkn.lastChar)]; decodedChar == sqltypes.DONTESCAPE {
				ch = tkn.lastChar
			} else {
				ch = uint16(decodedChar)
			}
			tkn.next()
		}
		if ch == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		buffer.WriteByte(byte(ch))
	}
	return typ, buffer.Bytes()
}

func (tkn *Tokenizer) scanCommentType1(prefix string) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteString(prefix)
	for tkn.lastChar != EOFCHAR {
		if tkn.lastChar == '\n' {
			tkn.ConsumeNext(buffer)
			break
		}
		tkn.ConsumeNext(buffer)
	}
	return COMMENT, buffer.Bytes()
}

func (tkn *Tokenizer) scanCommentType2() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteString("/*")
	for {
		if tkn.lastChar == '*' {
			tkn.ConsumeNext(buffer)
			if tkn.lastChar == '/' {
				tkn.ConsumeNext(buffer)
				break
			}
			continue
		}
		if tkn.lastChar == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		tkn.ConsumeNext(buffer)
	}
	return COMMENT, buffer.Bytes()
}

func (tkn *Tokenizer) ConsumeNext(buffer *bytes.Buffer) {
	if tkn.lastChar == EOFCHAR {
		// This should never happen.
		panic("unexpected EOF")
	}
	buffer.WriteByte(byte(tkn.lastChar))
	tkn.next()
}

func (tkn *Tokenizer) next() {
	if ch, err := tkn.InStream.ReadByte(); err != nil {
		// Only EOF is possible.
		tkn.lastChar = EOFCHAR
	} else {
		tkn.lastChar = uint16(ch)
	}
	tkn.Position++
}

func isLetter(ch uint16) bool {
	return 'a' <= ch && ch <= 'z' || 'A' <= ch && ch <= 'Z' || ch == '_' || ch == '@'
}

func digitVal(ch uint16) int {
	switch {
	case '0' <= ch && ch <= '9':
		return int(ch) - '0'
	case 'a' <= ch && ch <= 'f':
		return int(ch) - 'a' + 10
	case 'A' <= ch && ch <= 'F':
		return int(ch) - 'A' + 10
	}
	return 16 // larger than any legal digit val
}

func isDigit(ch uint16) bool {
	return '0' <= ch && ch <= '9'
}
<commit_msg>hex value support<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sqlparser

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/wandoulabs/cm/hack"
	"github.com/wandoulabs/cm/sqltypes"
)

const EOFCHAR = 0x100

// Tokenizer is the struct used to generate SQL
// tokens for the parser.
type Tokenizer struct {
	InStream      *strings.Reader
	AllowComments bool
	ForceEOF      bool
	lastChar      uint16
	Position      int
	errorToken    []byte
	LastError     string
	posVarIndex   int
	ParseTree     Statement
}

// NewStringTokenizer creates a new Tokenizer for the
// sql string.
func NewStringTokenizer(sql string) *Tokenizer {
	return &Tokenizer{InStream: strings.NewReader(sql)}
}

var keywords = map[string]int{
	"all":           ALL,
	"alter":         ALTER,
	"analyze":       ANALYZE,
	"and":           AND,
	"as":            AS,
	"asc":           ASC,
	"between":       BETWEEN,
	"by":            BY,
	"case":          CASE,
	"create":        CREATE,
	"cross":         CROSS,
	"default":       DEFAULT,
	"delete":        DELETE,
	"desc":          DESC,
	"describe":      DESCRIBE,
	"distinct":      DISTINCT,
	"drop":          DROP,
	"duplicate":     DUPLICATE,
	"else":          ELSE,
	"end":           END,
	"except":        EXCEPT,
	"exists":        EXISTS,
	"explain":       EXPLAIN,
	"for":           FOR,
	"force":         FORCE,
	"from":          FROM,
	"group":         GROUP,
	"having":        HAVING,
	"if":            IF,
	"ignore":        IGNORE,
	"in":            IN,
	"index":         INDEX,
	"inner":         INNER,
	"insert":        INSERT,
	"intersect":     INTERSECT,
	"into":          INTO,
	"is":            IS,
	"join":          JOIN,
	"key":           KEY,
	"left":          LEFT,
	"like":          LIKE,
	"limit":         LIMIT,
	"lock":          LOCK,
	"minus":         MINUS,
	"natural":       NATURAL,
	"not":           NOT,
	"null":          NULL,
	"on":            ON,
	"or":            OR,
	"order":         ORDER,
	"outer":         OUTER,
	"rename":        RENAME,
	"right":         RIGHT,
	"select":        SELECT,
	"set":           SET,
	"show":          SHOW,
	"straight_join": STRAIGHT_JOIN,
	"table":         TABLE,
	"then":          THEN,
	"to":            TO,
	"union":         UNION,
	"unique":        UNIQUE,
	"update":        UPDATE,
	"use":           USE,
	"using":         USING,
	"values":        VALUES,
	"view":          VIEW,
	"when":          WHEN,
	"where":         WHERE,
}

// Lex returns the next token form the Tokenizer.
// This function is used by go yacc.
func (tkn *Tokenizer) Lex(lval *yySymType) int {
	typ, val := tkn.Scan()
	for typ == COMMENT {
		if tkn.AllowComments {
			break
		}
		typ, val = tkn.Scan()
	}
	switch typ {
	case ID, STRING, NUMBER, VALUE_ARG, LIST_ARG, COMMENT:
		lval.bytes = val
	}
	tkn.errorToken = val
	return typ
}

// Error is called by go yacc if there's a parsing error.
func (tkn *Tokenizer) Error(err string) {
	buf := bytes.NewBuffer(make([]byte, 0, 32))
	if tkn.errorToken != nil {
		fmt.Fprintf(buf, "%s at position %v near %s", err, tkn.Position, tkn.errorToken)
	} else {
		fmt.Fprintf(buf, "%s at position %v", err, tkn.Position)
	}
	tkn.LastError = buf.String()
}

func (tkn *Tokenizer) scanHexValue(delim uint16, typ int) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	buffer.WriteString(`x'`)

	tkn.next()
	for {
		ch := tkn.lastChar
		tkn.next()
		if ch == delim {
			if tkn.lastChar == delim {
				tkn.next()
			} else {
				buffer.WriteByte(byte(ch))
				break
			}
		} else if ch == '\'' {
			if tkn.lastChar == EOFCHAR {
				return LEX_ERROR, buffer.Bytes()
			}
			if decodedChar := sqltypes.SqlDecodeMap[byte(tkn.lastChar)]; decodedChar == sqltypes.DONTESCAPE {
				ch = tkn.lastChar
			} else {
				ch = uint16(decodedChar)
			}
			tkn.next()
		}
		if ch == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		buffer.WriteByte(byte(ch))
	}
	return typ, buffer.Bytes()
}

// Scan scans the tokenizer for the next token and returns
// the token type and an optional value.
func (tkn *Tokenizer) Scan() (int, []byte) {
	if tkn.ForceEOF {
		return 0, nil
	}

	if tkn.lastChar == 0 {
		tkn.next()
	}
	tkn.skipBlank()
	switch ch := tkn.lastChar; {
	case isLetter(ch):
		if ch == 'x' {
			tkn.next()
			c := tkn.lastChar
			if c == '\'' {
				return tkn.scanHexValue('\'', STRING)
			}

			tkn.unReadByte()
		}
		return tkn.scanIdentifier()
	case isDigit(ch):
		return tkn.scanNumber(false)
	case ch == ':':
		return tkn.scanBindVar()
	default:
		tkn.next()
		switch ch {
		case EOFCHAR:
			return 0, nil
		case '=', ',', ';', '(', ')', '+', '*', '%', '&', '|', '^', '~':
			return int(ch), nil
		case '?':
			tkn.posVarIndex++
			buf := new(bytes.Buffer)
			fmt.Fprintf(buf, ":v%d", tkn.posVarIndex)
			return VALUE_ARG, buf.Bytes()
		case '.':
			if isDigit(tkn.lastChar) {
				return tkn.scanNumber(true)
			} else {
				return int(ch), nil
			}
		case '/':
			switch tkn.lastChar {
			case '/':
				tkn.next()
				return tkn.scanCommentType1("//")
			case '*':
				tkn.next()
				return tkn.scanCommentType2()
			default:
				return int(ch), nil
			}
		case '-':
			if tkn.lastChar == '-' {
				tkn.next()
				return tkn.scanCommentType1("--")
			} else {
				return int(ch), nil
			}
		case '<':
			switch tkn.lastChar {
			case '>':
				tkn.next()
				return NE, nil
			case '=':
				tkn.next()
				switch tkn.lastChar {
				case '>':
					tkn.next()
					return NULL_SAFE_EQUAL, nil
				default:
					return LE, nil
				}
			default:
				return int(ch), nil
			}
		case '>':
			if tkn.lastChar == '=' {
				tkn.next()
				return GE, nil
			} else {
				return int(ch), nil
			}
		case '!':
			if tkn.lastChar == '=' {
				tkn.next()
				return NE, nil
			} else {
				return LEX_ERROR, []byte("!")
			}
		case '\'', '"':
			return tkn.scanString(ch, STRING)
		case '`':
			return tkn.scanLiteralIdentifier()
		default:
			return LEX_ERROR, []byte{byte(ch)}
		}
	}
}

func (tkn *Tokenizer) skipBlank() {
	ch := tkn.lastChar
	for ch == ' ' || ch == '\n' || ch == '\r' || ch == '\t' {
		tkn.next()
		ch = tkn.lastChar
	}
}

func (tkn *Tokenizer) scanIdentifier() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	buffer.WriteByte(byte(tkn.lastChar))
	for tkn.next(); isLetter(tkn.lastChar) || isDigit(tkn.lastChar); tkn.next() {
		buffer.WriteByte(byte(tkn.lastChar))
	}
	lowered := bytes.ToLower(buffer.Bytes())
	if keywordId, found := keywords[hack.String(lowered)]; found {
		return keywordId, lowered
	}

	return ID, buffer.Bytes()
}

func (tkn *Tokenizer) scanLiteralIdentifier() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteByte(byte(tkn.lastChar))
	if !isLetter(tkn.lastChar) {
		return LEX_ERROR, buffer.Bytes()
	}
	for tkn.next(); isLetter(tkn.lastChar) || isDigit(tkn.lastChar); tkn.next() {
		buffer.WriteByte(byte(tkn.lastChar))
	}
	if tkn.lastChar != '`' {
		return LEX_ERROR, buffer.Bytes()
	}
	tkn.next()
	return ID, buffer.Bytes()
}

func (tkn *Tokenizer) scanBindVar() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteByte(byte(tkn.lastChar))
	token := VALUE_ARG
	tkn.next()
	if tkn.lastChar == ':' {
		token = LIST_ARG
		buffer.WriteByte(byte(tkn.lastChar))
		tkn.next()
	}
	if !isLetter(tkn.lastChar) {
		return LEX_ERROR, buffer.Bytes()
	}
	for isLetter(tkn.lastChar) || isDigit(tkn.lastChar) || tkn.lastChar == '.' {
		buffer.WriteByte(byte(tkn.lastChar))
		tkn.next()
	}
	return token, buffer.Bytes()
}

func (tkn *Tokenizer) scanMantissa(base int, buffer *bytes.Buffer) {
	for digitVal(tkn.lastChar) < base {
		tkn.ConsumeNext(buffer)
	}
}

func (tkn *Tokenizer) scanNumber(seenDecimalPoint bool) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	if seenDecimalPoint {
		buffer.WriteByte('.')
		tkn.scanMantissa(10, buffer)
		goto exponent
	}

	if tkn.lastChar == '0' {
		// int or float
		tkn.ConsumeNext(buffer)
		if tkn.lastChar == 'x' || tkn.lastChar == 'X' {
			// hexadecimal int
			tkn.ConsumeNext(buffer)
			tkn.scanMantissa(16, buffer)
		} else {
			// octal int or float
			seenDecimalDigit := false
			tkn.scanMantissa(8, buffer)
			if tkn.lastChar == '8' || tkn.lastChar == '9' {
				// illegal octal int or float
				seenDecimalDigit = true
				tkn.scanMantissa(10, buffer)
			}
			if tkn.lastChar == '.' || tkn.lastChar == 'e' || tkn.lastChar == 'E' {
				goto fraction
			}
			// octal int
			if seenDecimalDigit {
				return LEX_ERROR, buffer.Bytes()
			}
		}
		goto exit
	}

	// decimal int or float
	tkn.scanMantissa(10, buffer)

fraction:
	if tkn.lastChar == '.' {
		tkn.ConsumeNext(buffer)
		tkn.scanMantissa(10, buffer)
	}

exponent:
	if tkn.lastChar == 'e' || tkn.lastChar == 'E' {
		tkn.ConsumeNext(buffer)
		if tkn.lastChar == '+' || tkn.lastChar == '-' {
			tkn.ConsumeNext(buffer)
		}
		tkn.scanMantissa(10, buffer)
	}

exit:
	return NUMBER, buffer.Bytes()
}

func (tkn *Tokenizer) scanString(delim uint16, typ int) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	for {
		ch := tkn.lastChar
		tkn.next()
		if ch == delim {
			if tkn.lastChar == delim {
				tkn.next()
			} else {
				break
			}
		} else if ch == '\\' {
			if tkn.lastChar == EOFCHAR {
				return LEX_ERROR, buffer.Bytes()
			}
			if decodedChar := sqltypes.SqlDecodeMap[byte(tkn.lastChar)]; decodedChar == sqltypes.DONTESCAPE {
				ch = tkn.lastChar
			} else {
				ch = uint16(decodedChar)
			}
			tkn.next()
		}
		if ch == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		buffer.WriteByte(byte(ch))
	}
	return typ, buffer.Bytes()
}

func (tkn *Tokenizer) scanCommentType1(prefix string) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteString(prefix)
	for tkn.lastChar != EOFCHAR {
		if tkn.lastChar == '\n' {
			tkn.ConsumeNext(buffer)
			break
		}
		tkn.ConsumeNext(buffer)
	}
	return COMMENT, buffer.Bytes()
}

func (tkn *Tokenizer) scanCommentType2() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteString("/*")
	for {
		if tkn.lastChar == '*' {
			tkn.ConsumeNext(buffer)
			if tkn.lastChar == '/' {
				tkn.ConsumeNext(buffer)
				break
			}
			continue
		}
		if tkn.lastChar == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		tkn.ConsumeNext(buffer)
	}
	return COMMENT, buffer.Bytes()
}

func (tkn *Tokenizer) ConsumeNext(buffer *bytes.Buffer) {
	if tkn.lastChar == EOFCHAR {
		// This should never happen.
		panic("unexpected EOF")
	}
	buffer.WriteByte(byte(tkn.lastChar))
	tkn.next()
}

func (tkn *Tokenizer) unReadByte() {
	tkn.InStream.UnreadByte()
}

func (tkn *Tokenizer) next() {
	if ch, err := tkn.InStream.ReadByte(); err != nil {
		// Only EOF is possible.
		tkn.lastChar = EOFCHAR
	} else {
		tkn.lastChar = uint16(ch)
	}
	tkn.Position++
}

func isLetter(ch uint16) bool {
	return 'a' <= ch && ch <= 'z' || 'A' <= ch && ch <= 'Z' || ch == '_' || ch == '@'
}

func digitVal(ch uint16) int {
	switch {
	case '0' <= ch && ch <= '9':
		return int(ch) - '0'
	case 'a' <= ch && ch <= 'f':
		return int(ch) - 'a' + 10
	case 'A' <= ch && ch <= 'F':
		return int(ch) - 'A' + 10
	}
	return 16 // larger than any legal digit val
}

func isDigit(ch uint16) bool {
	return '0' <= ch && ch <= '9'
}
<|endoftext|>