<commit_before>package storage

import (
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/koding/logging"
)

var _ Interface = (*S3)(nil)

// S3 provides Storage functionality backed by S3.
type S3 struct {
	bucketName string
	s3         *s3.S3 // bucket `bucketName` holds the plans of terraform
	log        logging.Logger
}

// NewS3 creates a new S3 system.
func NewS3(key, secret, bucketName string, log logging.Logger) (*S3, error) {
	cfg := &aws.Config{
		Credentials: credentials.NewStaticCredentials(key, secret, ""),
		Region:      aws.String("us-east-1"), // we are only using us east
	}

	svc := s3.New(session.New(cfg))

	params := &s3.PutBucketAclInput{
		Bucket: aws.String(bucketName),
		ACL:    aws.String(s3.BucketCannedACLPrivate),
	}
	if _, err := svc.PutBucketAcl(params); err != nil {
		return nil, err
	}

	return &S3{
		bucketName: bucketName,
		s3:         svc,
		log:        log,
	}, nil
}

// BasePath returns bucket name.
func (s *S3) BasePath() (string, error) {
	return s.bucketName, nil
}

// Write writes to a s3 bucket.
func (s *S3) Write(path string, file io.Reader) error {
	params := &s3.PutObjectInput{
		ACL:         aws.String(s3.BucketCannedACLPrivate),
		Bucket:      aws.String(s.bucketName),
		Key:         aws.String(path),
		ContentType: aws.String("application/json"),
	}
	if rs, ok := file.(io.ReadSeeker); ok {
		params.Body = rs
		if size := s.guessSize(file); size != 0 {
			params.ContentLength = aws.Int64(size)
		}
	} else {
		// The file does not support seeking, thus can't be used
		// by the AWS streaming api - read it in-memory instead.
		content, err := ioutil.ReadAll(file)
		if err != nil {
			return err
		}
		params.Body = bytes.NewReader(content)
		params.ContentLength = aws.Int64(int64(len(content)))
	}

	_, err := s.s3.PutObject(params)
	// TODO(rjeczalik): make the write blocking with s3.WaitUntilObjectExists?
	return err
}

// Remove removes a file from a bucket.
func (s *S3) Remove(path string) error {
	params := &s3.DeleteObjectInput{
		Bucket: aws.String(s.bucketName),
		Key:    aws.String(path),
	}
	_, err := s.s3.DeleteObject(params)
	return err
}

// Read reads a file from a bucket.
//
// Caller is responsible for closing the request body.
func (s *S3) Read(path string) (io.Reader, error) {
	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucketName),
		Key:    aws.String(path),
	}
	resp, err := s.s3.GetObject(params)
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

// Clone clones the contents of a bucket to target storage
func (s *S3) Clone(path string, target Interface) error {
	params := &s3.ListObjectsInput{
		Bucket: aws.String(s.bucketName),
		// Prefix limits the response to keys that begin with the specified prefix.
		// You can use prefixes to separate a bucket into different groupings of keys.
		// (You can think of using prefix to make groups in the same way you'd use a
		// folder in a file system.)
		Prefix: nil,
		// Specifies the key to start with when listing objects in a bucket. Amazon
		// S3 returns object keys in alphabetical order, starting with key after the
		// marker in order.
		Marker: nil,
		// Sets the maximum number of keys returned in the response body. You can
		// add this to your request if you want to retrieve fewer than the default
		// 1000 keys.
		MaxKeys: nil,
	}
	if path != "" {
		// If path was an empty string, ListObject would successfuly return nothing.
		path = path + "/"
		params.Prefix = aws.String(path)
	}
	// read all elements in a bucket, we are gonna have more than 1000 items in
	// that bucket/folder, so we handle paginated response to copy them all.
	var targetErr error
	err := s.s3.ListObjectsPages(params, func(resp *s3.ListObjectsOutput, _ bool) bool {
		for _, obj := range resp.Contents {
			// write them all to target
			newPath := aws.StringValue(obj.Key)

			// if bucket is created but doesnt have any objects
			if newPath == path {
				continue
			}

			var r io.Reader
			r, targetErr = s.Read(newPath)
			if targetErr != nil {
				return false
			}

			targetErr = target.Write(newPath, r)

			// If the reader implements io.Closer interface, close the resource.
			s.ensureClosed(r, newPath)

			if targetErr != nil {
				return false
			}
		}
		return true
	})
	if targetErr != nil {
		return targetErr
	}
	return err
}

func (s *S3) ensureClosed(r io.Reader, path string) {
	if c, ok := r.(io.Closer); ok {
		if err := c.Close(); err != nil {
			s.log.Warning("failed closing resource path=%q: %s", path, err)
		}
	}
}

func (s *S3) guessSize(r io.Reader) int64 {
	// trying to be smart here:
	//
	//  https://github.com/golang/go/blob/0417872/src/net/http/request.go#L602-L609
	//
	switch v := r.(type) {
	case *bytes.Buffer:
		return int64(v.Len())
	case *bytes.Reader:
		return int64(v.Len())
	case *strings.Reader:
		return int64(v.Len())
	case *os.File:
		if fi, err := v.Stat(); err == nil {
			return fi.Size()
		}
	}
	s.log.Warning("unable to guess body size for %T", r)
	return 0
}
<commit_msg>terraformer/storage: simplify Write method<commit_after>package storage

import (
	"bytes"
	"io"
	"io/ioutil"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/koding/logging"
)

var _ Interface = (*S3)(nil)

// S3 provides Storage functionality backed by S3.
type S3 struct {
	bucketName string
	s3         *s3.S3 // bucket `bucketName` holds the plans of terraform
	log        logging.Logger
}

// NewS3 creates a new S3 system.
func NewS3(key, secret, bucketName string, log logging.Logger) (*S3, error) {
	cfg := &aws.Config{
		Credentials: credentials.NewStaticCredentials(key, secret, ""),
		Region:      aws.String("us-east-1"), // we are only using us east
	}

	svc := s3.New(session.New(cfg))

	params := &s3.PutBucketAclInput{
		Bucket: aws.String(bucketName),
		ACL:    aws.String(s3.BucketCannedACLPrivate),
	}
	if _, err := svc.PutBucketAcl(params); err != nil {
		return nil, err
	}

	return &S3{
		bucketName: bucketName,
		s3:         svc,
		log:        log,
	}, nil
}

// BasePath returns bucket name.
func (s *S3) BasePath() (string, error) {
	return s.bucketName, nil
}

// Write writes to a s3 bucket.
func (s *S3) Write(path string, file io.Reader) error {
	params := &s3.PutObjectInput{
		ACL:         aws.String(s3.BucketCannedACLPrivate),
		Bucket:      aws.String(s.bucketName),
		Key:         aws.String(path),
		ContentType: aws.String("application/json"),
	}
	body, ok := file.(io.ReadSeeker)
	if !ok {
		// The file does not support seeking, thus can't be used
		// by the AWS streaming api - read it in-memory instead.
		content, err := ioutil.ReadAll(file)
		if err != nil {
			return err
		}
		body = bytes.NewReader(content)
	}
	params.Body = body
	_, err := s.s3.PutObject(params)
	// TODO(rjeczalik): make the write blocking with s3.WaitUntilObjectExists?
	return err
}

// Remove removes a file from a bucket.
func (s *S3) Remove(path string) error {
	params := &s3.DeleteObjectInput{
		Bucket: aws.String(s.bucketName),
		Key:    aws.String(path),
	}
	_, err := s.s3.DeleteObject(params)
	return err
}

// Read reads a file from a bucket.
//
// Caller is responsible for closing the request body.
func (s *S3) Read(path string) (io.Reader, error) {
	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucketName),
		Key:    aws.String(path),
	}
	resp, err := s.s3.GetObject(params)
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

// Clone clones the contents of a bucket to target storage
func (s *S3) Clone(path string, target Interface) error {
	params := &s3.ListObjectsInput{
		Bucket: aws.String(s.bucketName),
		// Prefix limits the response to keys that begin with the specified prefix.
		// You can use prefixes to separate a bucket into different groupings of keys.
		// (You can think of using prefix to make groups in the same way you'd use a
		// folder in a file system.)
		Prefix: nil,
		// Specifies the key to start with when listing objects in a bucket. Amazon
		// S3 returns object keys in alphabetical order, starting with key after the
		// marker in order.
		Marker: nil,
		// Sets the maximum number of keys returned in the response body. You can
		// add this to your request if you want to retrieve fewer than the default
		// 1000 keys.
		MaxKeys: nil,
	}
	if path != "" {
		// If path was an empty string, ListObject would successfuly return nothing.
		path = path + "/"
		params.Prefix = aws.String(path)
	}
	// read all elements in a folder of the bucket, we are gonna have more than
	// 1000 items in it, so we handle paginated response to copy them all.
	var targetErr error
	err := s.s3.ListObjectsPages(params, func(resp *s3.ListObjectsOutput, _ bool) bool {
		for _, obj := range resp.Contents {
			// write them all to target
			newPath := aws.StringValue(obj.Key)

			// if bucket is created but doesnt have any objects
			if newPath == path {
				continue
			}

			var r io.Reader
			r, targetErr = s.Read(newPath)
			if targetErr != nil {
				return false
			}

			targetErr = target.Write(newPath, r)

			// If the reader implements io.Closer interface, close the resource.
			s.ensureClosed(r, newPath)

			if targetErr != nil {
				return false
			}
		}
		return true
	})
	if targetErr != nil {
		return targetErr
	}
	return err
}

func (s *S3) ensureClosed(r io.Reader, path string) {
	if c, ok := r.(io.Closer); ok {
		if err := c.Close(); err != nil {
			s.log.Warning("failed closing resource path=%q: %s", path, err)
		}
	}
}
<|endoftext|>


<commit_before>package models

import (
	"errors"
	"socialapi/request"
	"time"
	"github.com/koding/bongo"
)

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

	// Id of the channel
	ChannelId int64 `json:"channelId,string"     sql:"NOT NULL"`

	// Id of the message
	MessageId int64 `json:"messageId,string"     sql:"NOT NULL"`

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

	// Addition date of the message to the channel
	// this date will be update whever message added/removed/re-added to the channel
	AddedAt time.Time `json:"addedAt"            sql:"NOT NULL"`

	// Update time of the message/list
	RevisedAt time.Time `json:"revisedAt"        sql:"NOT NULL"`
}

func (c *ChannelMessageList) BeforeCreate() {
	c.AddedAt = time.Now()
	c.RevisedAt = time.Now()
}

func (c *ChannelMessageList) BeforeUpdate() {
	c.AddedAt = time.Now()
}

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

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

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

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

func (c ChannelMessageList) TableName() string {
	return "api.channel_message_list"
}

func NewChannelMessageList() *ChannelMessageList {
	return &ChannelMessageList{}
}

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

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

func (c *ChannelMessageList) Update() error {
	return bongo.B.Update(c)
}

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

func (c *ChannelMessageList) UnreadCount(cp *ChannelParticipant) (int, error) {
	if cp.ChannelId == 0 {
		return 0, errors.New("ChannelId is not set")
	}

	if cp.AccountId == 0 {
		return 0, errors.New("AccountId is not set")
	}

	if cp.LastSeenAt.IsZero() {
		return 0, errors.New("Last seen at date is not valid - it is zero")
	}

	return bongo.B.Count(c,
		"channel_id = ? and added_at > ?",
		cp.ChannelId,
		// todo change this format to get from a specific place
		cp.LastSeenAt.UTC().Format(time.RFC3339),
	)
}

func (c *ChannelMessageList) Create() error {
	return bongo.B.Create(c)
}

func (c *ChannelMessageList) CreateRaw() error {
	insertSql := "INSERT INTO " +
		c.TableName() +
		` ("channel_id","message_id","added_at","revised_at") VALUES ($1,$2,$3,$4) ` +
		"RETURNING ID"

	return bongo.B.DB.CommonDB().
		QueryRow(insertSql, c.ChannelId, c.MessageId, c.AddedAt, c.RevisedAt).
		Scan(&c.Id)
}

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

func (c *ChannelMessageList) List(q *request.Query, populateUnreadCount bool) (*HistoryResponse, error) {
	messageList, err := c.getMessages(q)
	if err != nil {
		return nil, err
	}

	if populateUnreadCount {
		messageList = c.populateUnreadCount(messageList)
	}

	hr := NewHistoryResponse()
	hr.MessageList = messageList
	return hr, nil
}

// populateUnreadCount adds unread count into message containers
func (c *ChannelMessageList) populateUnreadCount(messageList []*ChannelMessageContainer) []*ChannelMessageContainer {
	channel := NewChannel()
	channel.Id = c.ChannelId

	for i, message := range messageList {
		cml, err := channel.FetchMessageList(message.Message.Id)
		if err != nil {
			// helper.MustGetLogger().Error(err.Error())
			continue
		}

		count, err := NewMessageReply().UnreadCount(cml.MessageId, cml.RevisedAt)
		if err != nil {
			// helper.MustGetLogger().Error(err.Error())
			continue
		}
		messageList[i].UnreadRepliesCount = count
	}

	return messageList
}

func (c *ChannelMessageList) getMessages(q *request.Query) ([]*ChannelMessageContainer, error) {
	var messages []int64

	if c.ChannelId == 0 {
		return nil, errors.New("ChannelId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": c.ChannelId,
		},
		Pluck:      "message_id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
	}

	query.AddScope(SortedByAddedAt)
	query.AddScope(RemoveTrollContent(c, q.ShowExempt))

	bongoQuery := bongo.B.BuildQuery(c, query)
	if !q.From.IsZero() {
		bongoQuery = bongoQuery.Where("added_at < ?", q.From)
	}

	if err := bongo.CheckErr(
		bongoQuery.Pluck(query.Pluck, &messages),
	); err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessages, err := parent.FetchByIds(messages)
	if err != nil {
		return nil, err
	}

	populatedChannelMessages, err := c.populateChannelMessages(channelMessages, q)
	if err != nil {
		return nil, err
	}

	return populatedChannelMessages, nil
}

func (c *ChannelMessageList) IsInChannel(messageId, channelId int64) (bool, error) {
	if messageId == 0 || channelId == 0 {
		return false, errors.New("channelId/messageId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channelId,
			"message_id": messageId,
		},
	}

	err := c.One(query)
	if err == nil {
		return true, nil
	}

	if err == bongo.RecordNotFound {
		return false, nil
	}

	return false, err
}

func (c *ChannelMessageList) populateChannelMessages(channelMessages []ChannelMessage, query *request.Query) ([]*ChannelMessageContainer, error) {
	channelMessageCount := len(channelMessages)

	populatedChannelMessages := make([]*ChannelMessageContainer, channelMessageCount)

	if channelMessageCount == 0 {
		return populatedChannelMessages, nil
	}

	for i := 0; i < channelMessageCount; i++ {
		cm := channelMessages[i]
		cmc, err := cm.BuildMessage(query)
		if err != nil {
			return nil, err
		}

		populatedChannelMessages[i] = cmc
	}

	return populatedChannelMessages, nil
}

func (c *ChannelMessageList) FetchMessageChannelIds(messageId int64) ([]int64, error) {
	var channelIds []int64

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": messageId,
		},
		Pluck: "channel_id",
	}

	err := bongo.B.Some(c, &channelIds, q)
	if err != nil {
		return nil, err
	}

	return channelIds, nil
}

func (c *ChannelMessageList) FetchMessageChannels(messageId int64) ([]Channel, error) {
	channelIds, err := c.FetchMessageChannelIds(messageId)
	if err != nil {
		return nil, err
	}

	return NewChannel().FetchByIds(channelIds)
}

func (c *ChannelMessageList) FetchMessageIdsByChannelId(channelId int64, q *request.Query) ([]int64, error) {
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channelId,
		},
		Pluck:      "message_id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort: map[string]string{
			"added_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
}

// seperate this fucntion into modelhelper
// as setting it to a variadic function
func (c *ChannelMessageList) DeleteMessagesBySelector(selector map[string]interface{}) error {
	var cmls []ChannelMessageList

	err := bongo.B.Some(c, &cmls, &bongo.Query{Selector: selector})
	if err != nil {
		return err
	}

	for _, cml := range cmls {
		if err := cml.Delete(); err != nil {
			return err
		}
	}
	return nil
}

func (c *ChannelMessageList) UpdateAddedAt(channelId, messageId int64) error {
	if messageId == 0 || channelId == 0 {
		return errors.New("channelId/messageId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channelId,
			"message_id": messageId,
		},
	}

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

	c.AddedAt = time.Now().UTC()
	return c.Update()
}
<commit_msg>Social: check and mark channel message list before<commit_after>package models

import (
	"errors"
	"socialapi/request"
	"time"

	"github.com/koding/bongo"
)

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

	// Id of the channel
	ChannelId int64 `json:"channelId,string"     sql:"NOT NULL"`

	// Id of the message
	MessageId int64 `json:"messageId,string"     sql:"NOT NULL"`

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

	// Addition date of the message to the channel
	// this date will be update whever message added/removed/re-added to the channel
	AddedAt time.Time `json:"addedAt"            sql:"NOT NULL"`

	// Update time of the message/list
	RevisedAt time.Time `json:"revisedAt"        sql:"NOT NULL"`
}

func (c *ChannelMessageList) BeforeCreate() error {
	c.AddedAt = time.Now()
	c.RevisedAt = time.Now()
	return c.MarkIfExempt()
}

func (c *ChannelMessageList) BeforeUpdate() error {
	c.AddedAt = time.Now()
	return c.MarkIfExempt()
}

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

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

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

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

func (c ChannelMessageList) TableName() string {
	return "api.channel_message_list"
}

func NewChannelMessageList() *ChannelMessageList {
	return &ChannelMessageList{}
}

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

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

func (c *ChannelMessageList) Update() error {
	return bongo.B.Update(c)
}

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

func (c *ChannelMessageList) UnreadCount(cp *ChannelParticipant) (int, error) {
	if cp.ChannelId == 0 {
		return 0, errors.New("ChannelId is not set")
	}

	if cp.AccountId == 0 {
		return 0, errors.New("AccountId is not set")
	}

	if cp.LastSeenAt.IsZero() {
		return 0, errors.New("Last seen at date is not valid - it is zero")
	}

	return bongo.B.Count(c,
		"channel_id = ? and added_at > ?",
		cp.ChannelId,
		// todo change this format to get from a specific place
		cp.LastSeenAt.UTC().Format(time.RFC3339),
	)
}

func (c *ChannelMessageList) Create() error {
	return bongo.B.Create(c)
}

func (c *ChannelMessageList) CreateRaw() error {
	insertSql := "INSERT INTO " +
		c.TableName() +
		` ("channel_id","message_id","added_at","revised_at") VALUES ($1,$2,$3,$4) ` +
		"RETURNING ID"

	return bongo.B.DB.CommonDB().
		QueryRow(insertSql, c.ChannelId, c.MessageId, c.AddedAt, c.RevisedAt).
		Scan(&c.Id)
}

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

func (c *ChannelMessageList) List(q *request.Query, populateUnreadCount bool) (*HistoryResponse, error) {
	messageList, err := c.getMessages(q)
	if err != nil {
		return nil, err
	}

	if populateUnreadCount {
		messageList = c.populateUnreadCount(messageList)
	}

	hr := NewHistoryResponse()
	hr.MessageList = messageList
	return hr, nil
}

// populateUnreadCount adds unread count into message containers
func (c *ChannelMessageList) populateUnreadCount(messageList []*ChannelMessageContainer) []*ChannelMessageContainer {
	channel := NewChannel()
	channel.Id = c.ChannelId

	for i, message := range messageList {
		cml, err := channel.FetchMessageList(message.Message.Id)
		if err != nil {
			// helper.MustGetLogger().Error(err.Error())
			continue
		}

		count, err := NewMessageReply().UnreadCount(cml.MessageId, cml.RevisedAt)
		if err != nil {
			// helper.MustGetLogger().Error(err.Error())
			continue
		}
		messageList[i].UnreadRepliesCount = count
	}

	return messageList
}

func (c *ChannelMessageList) getMessages(q *request.Query) ([]*ChannelMessageContainer, error) {
	var messages []int64

	if c.ChannelId == 0 {
		return nil, errors.New("ChannelId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": c.ChannelId,
		},
		Pluck:      "message_id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
	}

	query.AddScope(SortedByAddedAt)
	query.AddScope(RemoveTrollContent(c, q.ShowExempt))

	bongoQuery := bongo.B.BuildQuery(c, query)
	if !q.From.IsZero() {
		bongoQuery = bongoQuery.Where("added_at < ?", q.From)
	}

	if err := bongo.CheckErr(
		bongoQuery.Pluck(query.Pluck, &messages),
	); err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessages, err := parent.FetchByIds(messages)
	if err != nil {
		return nil, err
	}

	populatedChannelMessages, err := c.populateChannelMessages(channelMessages, q)
	if err != nil {
		return nil, err
	}

	return populatedChannelMessages, nil
}

func (c *ChannelMessageList) IsInChannel(messageId, channelId int64) (bool, error) {
	if messageId == 0 || channelId == 0 {
		return false, errors.New("channelId/messageId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channelId,
			"message_id": messageId,
		},
	}

	err := c.One(query)
	if err == nil {
		return true, nil
	}

	if err == bongo.RecordNotFound {
		return false, nil
	}

	return false, err
}

func (c *ChannelMessageList) populateChannelMessages(channelMessages []ChannelMessage, query *request.Query) ([]*ChannelMessageContainer, error) {
	channelMessageCount := len(channelMessages)

	populatedChannelMessages := make([]*ChannelMessageContainer, channelMessageCount)

	if channelMessageCount == 0 {
		return populatedChannelMessages, nil
	}

	for i := 0; i < channelMessageCount; i++ {
		cm := channelMessages[i]
		cmc, err := cm.BuildMessage(query)
		if err != nil {
			return nil, err
		}

		populatedChannelMessages[i] = cmc
	}

	return populatedChannelMessages, nil
}

func (c *ChannelMessageList) FetchMessageChannelIds(messageId int64) ([]int64, error) {
	var channelIds []int64

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": messageId,
		},
		Pluck: "channel_id",
	}

	err := bongo.B.Some(c, &channelIds, q)
	if err != nil {
		return nil, err
	}

	return channelIds, nil
}

func (c *ChannelMessageList) FetchMessageChannels(messageId int64) ([]Channel, error) {
	channelIds, err := c.FetchMessageChannelIds(messageId)
	if err != nil {
		return nil, err
	}

	return NewChannel().FetchByIds(channelIds)
}

func (c *ChannelMessageList) FetchMessageIdsByChannelId(channelId int64, q *request.Query) ([]int64, error) {
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channelId,
		},
		Pluck:      "message_id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort: map[string]string{
			"added_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
}

// seperate this fucntion into modelhelper
// as setting it to a variadic function
func (c *ChannelMessageList) DeleteMessagesBySelector(selector map[string]interface{}) error {
	var cmls []ChannelMessageList

	err := bongo.B.Some(c, &cmls, &bongo.Query{Selector: selector})
	if err != nil {
		return err
	}

	for _, cml := range cmls {
		if err := cml.Delete(); err != nil {
			return err
		}
	}
	return nil
}

func (c *ChannelMessageList) UpdateAddedAt(channelId, messageId int64) error {
	if messageId == 0 || channelId == 0 {
		return errors.New("channelId/messageId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channelId,
			"message_id": messageId,
		},
	}

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

	c.AddedAt = time.Now().UTC()
	return c.Update()
}

func (c *ChannelMessageList) MarkIfExempt() error {
	isExempt, err := c.isExempt()
	if err != nil {
		return err
	}

	if isExempt {
		c.MetaBits.MarkTroll()
	}

	return nil
}

func (c *ChannelMessageList) isExempt() (bool, error) {
	// return early if channel is already exempt
	if c.MetaBits.IsTroll() {
		return true, nil
	}

	if c.MessageId == 0 {
		return false, errors.New("message id is not set for exempt check")
	}

	cm := NewChannelMessage()
	cm.Id = c.MessageId

	return cm.isExempt()
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"strings"

	"github.com/fsouza/go-dockerclient"
	"github.com/gin-gonic/gin"
)

func stopContainer(c *gin.Context) {
	var container struct {
		ID string `json:"id"`
	}

	//reflect.ValueOf(&t).MethodByName("Foo").Call([]reflect.Value{})
	endpoint := "unix:///var/run/docker.sock"
	client, _ := docker.NewClient(endpoint)

	err := client.StopContainer(container.ID, 1)
	if err != nil {
		e := ErrorResponse{}
		e.ErrorMessage = "error stopping containers:" + err.Error()
		c.JSON(http.StatusInternalServerError, e)
		return
	}
	m := Message{}
	m.Message = "continer stopped successfully"
	c.JSON(http.StatusOK, m)
}

func stopAllContainers(c *gin.Context) {

	var stopErrors [][]string
	for i := range stopErrors {
		stopErrors[i] = make([]string, 0)
	}
	//reflect.ValueOf(&t).MethodByName("Foo").Call([]reflect.Value{})
	endpoint := "unix:///var/run/docker.sock"
	client, _ := docker.NewClient(endpoint)
	runningContainers, err := client.ListContainers(docker.ListContainersOptions{All: false})

	if err != nil {
		e := ErrorResponse{}
		e.ErrorMessage = "error getting all containers:" + err.Error()
		c.JSON(http.StatusInternalServerError, e)
		return
	}

	for _, v := range runningContainers {
		err := client.StopContainer(v.ID, 1)
		if err != nil {
			stopErrors = append(stopErrors, v.Names)
		}
	}

	if len(stopErrors) != 0 {
		var errCon string
		for _, v := range stopErrors {
			errCon += "[" + strings.Join(v, "|") + "]"
		}
		e := ErrorResponse{}
		e.ErrorMessage = "error stopping containers:" + errCon
		c.JSON(http.StatusInternalServerError, e)
		return
	}

	m := Message{}
	m.Message = "all containers successfully stopped"
	c.JSON(http.StatusOK, m)
}

func inspectContainer(c *gin.Context) {
	cID := c.Param("id")
	endpoint := "unix:///var/run/docker.sock"
	client, _ := docker.NewClient(endpoint)
	con, err := client.InspectContainer(cID)
	if err != nil {
		e := ErrorResponse{"error while trying to inspect container:" + err.Error()}
		c.JSON(http.StatusInternalServerError, e)
	}

	c.JSON(http.StatusOK, con)
}
<commit_msg>Fixed Agent not parsing the request body. Now, stopping a container works as intended.<commit_after>package main

import (
	"net/http"
	"strings"

	"github.com/fsouza/go-dockerclient"
	"github.com/gin-gonic/gin"
)

func stopContainer(c *gin.Context) {
	var container struct {
		ID string `json:"id"`
	}

	err := c.BindJSON(&container)
	if err != nil {
		e := ErrorResponse{}
		e.ErrorMessage = "error parsing json body:" + err.Error()
		c.JSON(http.StatusInternalServerError, e)
		return
	}
	//reflect.ValueOf(&t).MethodByName("Foo").Call([]reflect.Value{})
	endpoint := "unix:///var/run/docker.sock"
	client, _ := docker.NewClient(endpoint)

	err = client.StopContainer(container.ID, 10)
	if err != nil {
		e := ErrorResponse{}
		e.ErrorMessage = "error stopping containers:" + err.Error()
		c.JSON(http.StatusInternalServerError, e)
		return
	}
	m := Message{}
	m.Message = "continer stopped successfully"
	c.JSON(http.StatusOK, m)
}

func stopAllContainers(c *gin.Context) {

	var stopErrors [][]string
	for i := range stopErrors {
		stopErrors[i] = make([]string, 0)
	}
	//reflect.ValueOf(&t).MethodByName("Foo").Call([]reflect.Value{})
	endpoint := "unix:///var/run/docker.sock"
	client, _ := docker.NewClient(endpoint)
	runningContainers, err := client.ListContainers(docker.ListContainersOptions{All: false})

	if err != nil {
		e := ErrorResponse{}
		e.ErrorMessage = "error getting all containers:" + err.Error()
		c.JSON(http.StatusInternalServerError, e)
		return
	}

	for _, v := range runningContainers {
		err := client.StopContainer(v.ID, 1)
		if err != nil {
			stopErrors = append(stopErrors, v.Names)
		}
	}

	if len(stopErrors) != 0 {
		var errCon string
		for _, v := range stopErrors {
			errCon += "[" + strings.Join(v, "|") + "]"
		}
		e := ErrorResponse{}
		e.ErrorMessage = "error stopping containers:" + errCon
		c.JSON(http.StatusInternalServerError, e)
		return
	}

	m := Message{}
	m.Message = "all containers successfully stopped"
	c.JSON(http.StatusOK, m)
}

func inspectContainer(c *gin.Context) {
	cID := c.Param("id")
	endpoint := "unix:///var/run/docker.sock"
	client, _ := docker.NewClient(endpoint)
	con, err := client.InspectContainer(cID)
	if err != nil {
		e := ErrorResponse{"error while trying to inspect container:" + err.Error()}
		c.JSON(http.StatusInternalServerError, e)
	}

	c.JSON(http.StatusOK, con)
}
<|endoftext|>


<commit_before>package testutils

import (
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"google.golang.org/grpc"

	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/ca"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

// TestExecutor is executor for integration tests
type TestExecutor struct {
	mu   sync.Mutex
	desc *api.NodeDescription
}

// Describe just returns empty NodeDescription.
func (e *TestExecutor) Describe(ctx context.Context) (*api.NodeDescription, error) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.desc == nil {
		return &api.NodeDescription{}, nil
	}
	return e.desc.Copy(), nil
}

// Configure does nothing.
func (e *TestExecutor) Configure(ctx context.Context, node *api.Node) error {
	return nil
}

// SetNetworkBootstrapKeys does nothing.
func (e *TestExecutor) SetNetworkBootstrapKeys([]*api.EncryptionKey) error {
	return nil
}

// Controller returns TestController.
func (e *TestExecutor) Controller(t *api.Task) (exec.Controller, error) {
	return &TestController{
		ch: make(chan struct{}),
	}, nil
}

// UpdateNodeDescription sets the node description on the test executor
func (e *TestExecutor) UpdateNodeDescription(newDesc *api.NodeDescription) {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.desc = newDesc
}

// TestController is dummy channel based controller for tests.
type TestController struct {
	ch        chan struct{}
	closeOnce sync.Once
}

// Update does nothing.
func (t *TestController) Update(ctx context.Context, task *api.Task) error {
	return nil
}

// Prepare does nothing.
func (t *TestController) Prepare(ctx context.Context) error {
	return nil
}

// Start does nothing.
func (t *TestController) Start(ctx context.Context) error {
	return nil
}

// Wait waits on internal channel.
func (t *TestController) Wait(ctx context.Context) error {
	select {
	case <-t.ch:
	case <-ctx.Done():
	}
	return nil
}

// Shutdown closes internal channel
func (t *TestController) Shutdown(ctx context.Context) error {
	t.closeOnce.Do(func() {
		close(t.ch)
	})
	return nil
}

// Terminate closes internal channel if it wasn't closed before.
func (t *TestController) Terminate(ctx context.Context) error {
	t.closeOnce.Do(func() {
		close(t.ch)
	})
	return nil
}

// Remove does nothing.
func (t *TestController) Remove(ctx context.Context) error {
	return nil
}

// Close does nothing.
func (t *TestController) Close() error {
	t.closeOnce.Do(func() {
		close(t.ch)
	})
	return nil
}

// SessionHandler is an injectable function that can be used handle session requests
type SessionHandler func(*api.SessionRequest, api.Dispatcher_SessionServer) error

// MockDispatcher is a fake dispatcher that one agent at a time can connect to
type MockDispatcher struct {
	mu             sync.Mutex
	sessionCh      chan *api.SessionMessage
	openSession    *api.SessionRequest
	closedSessions []*api.SessionRequest
	sessionHandler SessionHandler

	Addr string
}

// UpdateTaskStatus is not implemented
func (m *MockDispatcher) UpdateTaskStatus(context.Context, *api.UpdateTaskStatusRequest) (*api.UpdateTaskStatusResponse, error) {
	panic("not implemented")
}

// Tasks keeps an open stream until canceled
func (m *MockDispatcher) Tasks(_ *api.TasksRequest, stream api.Dispatcher_TasksServer) error {
	select {
	case <-stream.Context().Done():
	}
	return nil
}

// Assignments keeps an open stream until canceled
func (m *MockDispatcher) Assignments(_ *api.AssignmentsRequest, stream api.Dispatcher_AssignmentsServer) error {
	select {
	case <-stream.Context().Done():
	}
	return nil
}

// Heartbeat always successfully heartbeats
func (m *MockDispatcher) Heartbeat(context.Context, *api.HeartbeatRequest) (*api.HeartbeatResponse, error) {
	return &api.HeartbeatResponse{Period: time.Second * 5}, nil
}

// Session allows a session to be established, and sends the node info
func (m *MockDispatcher) Session(r *api.SessionRequest, stream api.Dispatcher_SessionServer) error {
	m.mu.Lock()
	handler := m.sessionHandler
	m.openSession = r
	m.mu.Unlock()
	defer func() {
		m.mu.Lock()
		defer m.mu.Unlock()
		m.closedSessions = append(m.closedSessions, m.openSession)
		m.openSession = nil
	}()

	if handler != nil {
		return handler(r, stream)
	}

	// send the initial message first
	if err := stream.Send(&api.SessionMessage{
		SessionID: r.SessionID,
		Managers: []*api.WeightedPeer{
			{
				Peer: &api.Peer{Addr: m.Addr},
			},
		},
	}); err != nil {
		return err
	}

	ctx := stream.Context()
	for {
		select {
		case msg := <-m.sessionCh:
			msg.SessionID = r.SessionID
			if err := stream.Send(msg); err != nil {
				return err
			}
		case <-ctx.Done():
			return nil
		}
	}
}

// GetSessions return all the established and closed sessions
func (m *MockDispatcher) GetSessions() (*api.SessionRequest, []*api.SessionRequest) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.openSession, m.closedSessions
}

// SessionMessageChannel returns a writable channel to inject session messages
func (m *MockDispatcher) SessionMessageChannel() chan<- *api.SessionMessage {
	return m.sessionCh
}

// SetSessionHandler lets you inject a custom function to handle session requests
func (m *MockDispatcher) SetSessionHandler(s SessionHandler) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.sessionHandler = s
}

// NewMockDispatcher starts and returns a mock dispatcher instance that can be connected to
func NewMockDispatcher(t *testing.T, secConfig *ca.SecurityConfig, local bool) (*MockDispatcher, func()) {
	var (
		l       net.Listener
		err     error
		addr    string
		cleanup func()
	)
	if local {
		tempDir, err := ioutil.TempDir("", "local-dispatcher-socket")
		require.NoError(t, err)
		addr = filepath.Join(tempDir, "socket")
		l, err = net.Listen("unix", addr)
		require.NoError(t, err)
		cleanup = func() {
			os.RemoveAll(tempDir)
		}
	} else {
		l, err = net.Listen("tcp", "127.0.0.1:0")
		require.NoError(t, err)
		addr = l.Addr().String()
	}

	serverOpts := []grpc.ServerOption{grpc.Creds(secConfig.ServerTLSCreds)}
	s := grpc.NewServer(serverOpts...)

	m := &MockDispatcher{
		Addr:      addr,
		sessionCh: make(chan *api.SessionMessage, 1),
	}
	api.RegisterDispatcherServer(s, m)
	go s.Serve(l)
	return m, func() {
		l.Close()
		s.Stop()
		if cleanup != nil {
			cleanup()
		}
	}
}
<commit_msg>It's possible for a new session to start before the previous session has completely finished cleaning up in the fake dispatcher.  Therefore, in the cleanup, don't set the current open session to nil without checking if the current open session has changed.  This should prevent some false negative test assertions where it seems like a session has never been established even though it has.<commit_after>package testutils

import (
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"google.golang.org/grpc"

	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/ca"
	"github.com/docker/swarmkit/identity"
	"github.com/docker/swarmkit/log"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

// TestExecutor is executor for integration tests
type TestExecutor struct {
	mu   sync.Mutex
	desc *api.NodeDescription
}

// Describe just returns empty NodeDescription.
func (e *TestExecutor) Describe(ctx context.Context) (*api.NodeDescription, error) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.desc == nil {
		return &api.NodeDescription{}, nil
	}
	return e.desc.Copy(), nil
}

// Configure does nothing.
func (e *TestExecutor) Configure(ctx context.Context, node *api.Node) error {
	return nil
}

// SetNetworkBootstrapKeys does nothing.
func (e *TestExecutor) SetNetworkBootstrapKeys([]*api.EncryptionKey) error {
	return nil
}

// Controller returns TestController.
func (e *TestExecutor) Controller(t *api.Task) (exec.Controller, error) {
	return &TestController{
		ch: make(chan struct{}),
	}, nil
}

// UpdateNodeDescription sets the node description on the test executor
func (e *TestExecutor) UpdateNodeDescription(newDesc *api.NodeDescription) {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.desc = newDesc
}

// TestController is dummy channel based controller for tests.
type TestController struct {
	ch        chan struct{}
	closeOnce sync.Once
}

// Update does nothing.
func (t *TestController) Update(ctx context.Context, task *api.Task) error {
	return nil
}

// Prepare does nothing.
func (t *TestController) Prepare(ctx context.Context) error {
	return nil
}

// Start does nothing.
func (t *TestController) Start(ctx context.Context) error {
	return nil
}

// Wait waits on internal channel.
func (t *TestController) Wait(ctx context.Context) error {
	select {
	case <-t.ch:
	case <-ctx.Done():
	}
	return nil
}

// Shutdown closes internal channel
func (t *TestController) Shutdown(ctx context.Context) error {
	t.closeOnce.Do(func() {
		close(t.ch)
	})
	return nil
}

// Terminate closes internal channel if it wasn't closed before.
func (t *TestController) Terminate(ctx context.Context) error {
	t.closeOnce.Do(func() {
		close(t.ch)
	})
	return nil
}

// Remove does nothing.
func (t *TestController) Remove(ctx context.Context) error {
	return nil
}

// Close does nothing.
func (t *TestController) Close() error {
	t.closeOnce.Do(func() {
		close(t.ch)
	})
	return nil
}

// SessionHandler is an injectable function that can be used handle session requests
type SessionHandler func(*api.SessionRequest, api.Dispatcher_SessionServer) error

// MockDispatcher is a fake dispatcher that one agent at a time can connect to
type MockDispatcher struct {
	mu             sync.Mutex
	sessionCh      chan *api.SessionMessage
	openSession    *api.SessionRequest
	closedSessions []*api.SessionRequest
	sessionHandler SessionHandler

	Addr string
}

// UpdateTaskStatus is not implemented
func (m *MockDispatcher) UpdateTaskStatus(context.Context, *api.UpdateTaskStatusRequest) (*api.UpdateTaskStatusResponse, error) {
	panic("not implemented")
}

// Tasks keeps an open stream until canceled
func (m *MockDispatcher) Tasks(_ *api.TasksRequest, stream api.Dispatcher_TasksServer) error {
	select {
	case <-stream.Context().Done():
	}
	return nil
}

// Assignments keeps an open stream until canceled
func (m *MockDispatcher) Assignments(_ *api.AssignmentsRequest, stream api.Dispatcher_AssignmentsServer) error {
	select {
	case <-stream.Context().Done():
	}
	return nil
}

// Heartbeat always successfully heartbeats
func (m *MockDispatcher) Heartbeat(context.Context, *api.HeartbeatRequest) (*api.HeartbeatResponse, error) {
	return &api.HeartbeatResponse{Period: time.Second * 5}, nil
}

// Session allows a session to be established, and sends the node info
func (m *MockDispatcher) Session(r *api.SessionRequest, stream api.Dispatcher_SessionServer) error {
	m.mu.Lock()
	handler := m.sessionHandler
	m.openSession = r
	m.mu.Unlock()
	sessionID := identity.NewID()

	defer func() {
		m.mu.Lock()
		defer m.mu.Unlock()
		log.G(stream.Context()).Debugf("non-dispatcher side closed session: %s", sessionID)
		m.closedSessions = append(m.closedSessions, r)
		if m.openSession == r { // only overwrite session if it hasn't changed
			m.openSession = nil
		}
	}()

	if handler != nil {
		return handler(r, stream)
	}

	// send the initial message first
	if err := stream.Send(&api.SessionMessage{
		SessionID: sessionID,
		Managers: []*api.WeightedPeer{
			{
				Peer: &api.Peer{Addr: m.Addr},
			},
		},
	}); err != nil {
		return err
	}

	ctx := stream.Context()
	for {
		select {
		case msg := <-m.sessionCh:
			msg.SessionID = sessionID
			if err := stream.Send(msg); err != nil {
				return err
			}
		case <-ctx.Done():
			return nil
		}
	}
}

// GetSessions return all the established and closed sessions
func (m *MockDispatcher) GetSessions() (*api.SessionRequest, []*api.SessionRequest) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.openSession, m.closedSessions
}

// SessionMessageChannel returns a writable channel to inject session messages
func (m *MockDispatcher) SessionMessageChannel() chan<- *api.SessionMessage {
	return m.sessionCh
}

// SetSessionHandler lets you inject a custom function to handle session requests
func (m *MockDispatcher) SetSessionHandler(s SessionHandler) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.sessionHandler = s
}

// NewMockDispatcher starts and returns a mock dispatcher instance that can be connected to
func NewMockDispatcher(t *testing.T, secConfig *ca.SecurityConfig, local bool) (*MockDispatcher, func()) {
	var (
		l       net.Listener
		err     error
		addr    string
		cleanup func()
	)
	if local {
		tempDir, err := ioutil.TempDir("", "local-dispatcher-socket")
		require.NoError(t, err)
		addr = filepath.Join(tempDir, "socket")
		l, err = net.Listen("unix", addr)
		require.NoError(t, err)
		cleanup = func() {
			os.RemoveAll(tempDir)
		}
	} else {
		l, err = net.Listen("tcp", "127.0.0.1:0")
		require.NoError(t, err)
		addr = l.Addr().String()
	}

	serverOpts := []grpc.ServerOption{grpc.Creds(secConfig.ServerTLSCreds)}
	s := grpc.NewServer(serverOpts...)

	m := &MockDispatcher{
		Addr:      addr,
		sessionCh: make(chan *api.SessionMessage, 1),
	}
	api.RegisterDispatcherServer(s, m)
	go s.Serve(l)
	return m, func() {
		l.Close()
		s.Stop()
		if cleanup != nil {
			cleanup()
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2012 - Cloud Instruments Co., Ltd.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package seelog

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
)

// fileWriter is used to write to a file.
type fileWriter struct {
	innerWriter io.WriteCloser
	fileName    string
}

// Creates a new file and a corresponding writer. Returns error, if the file couldn't be created.
func newFileWriter(fileName string) (writer *fileWriter, err error) {
	newWriter := new(fileWriter)
	newWriter.fileName = fileName

	return newWriter, nil
}

func (fw *fileWriter) Close() error {
	if fw.innerWriter != nil {
		err := fw.innerWriter.Close()
		if err != nil {
			return err
		}
		fw.innerWriter = nil
	}
	return nil
}

// Create folder and file on WriteLog/Write first call
func (fw *fileWriter) Write(bytes []byte) (n int, err error) {
	if fw.innerWriter == nil {
		if err := fw.createFile(); err != nil {
			return 0, err
		}
	}
	return fw.innerWriter.Write(bytes)
}

func (fw *fileWriter) createFile() error {
	folder, _ := filepath.Split(fw.fileName)
	var err error

	if 0 != len(folder) {
		err = os.MkdirAll(folder, defaultDirectoryPermissions)
		if err != nil {
			return err
		}
	}

	// If exists
	_, err = os.Lstat(fw.fileName)
	if nil == err {
		fw.innerWriter, err = os.OpenFile(fw.fileName, os.O_WRONLY|os.O_APPEND, defaultFilePermissions)
	} else {
		fw.innerWriter, err = os.Create(fw.fileName)
	}

	if err != nil {
		return err
	}

	return nil
}

func (fw *fileWriter) String() string {
	return fmt.Sprintf("File writer: %s", fw.fileName)
}
<commit_msg>Avoid file creation\opening race in file writer.<commit_after>// Copyright (c) 2012 - Cloud Instruments Co., Ltd.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package seelog

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
)

// fileWriter is used to write to a file.
type fileWriter struct {
	innerWriter io.WriteCloser
	fileName    string
}

// Creates a new file and a corresponding writer. Returns error, if the file couldn't be created.
func newFileWriter(fileName string) (writer *fileWriter, err error) {
	newWriter := new(fileWriter)
	newWriter.fileName = fileName

	return newWriter, nil
}

func (fw *fileWriter) Close() error {
	if fw.innerWriter != nil {
		err := fw.innerWriter.Close()
		if err != nil {
			return err
		}
		fw.innerWriter = nil
	}
	return nil
}

// Create folder and file on WriteLog/Write first call
func (fw *fileWriter) Write(bytes []byte) (n int, err error) {
	if fw.innerWriter == nil {
		if err := fw.createFile(); err != nil {
			return 0, err
		}
	}
	return fw.innerWriter.Write(bytes)
}

func (fw *fileWriter) createFile() error {
	folder, _ := filepath.Split(fw.fileName)
	var err error

	if 0 != len(folder) {
		err = os.MkdirAll(folder, defaultDirectoryPermissions)
		if err != nil {
			return err
		}
	}

	// If exists
	fw.innerWriter, err = os.OpenFile(fw.fileName, os.O_WRONLY|os.O_APPEND|os.O_CREATE, defaultFilePermissions)

	if err != nil {
		return err
	}

	return nil
}

func (fw *fileWriter) String() string {
	return fmt.Sprintf("File writer: %s", fw.fileName)
}
<|endoftext|>


<commit_before>package dgrl

import (
	"strings"
)

type Node interface {
	Parent() *Branch
	SetParent(*Branch)
	Type() int
	Key() string
	SetKey(key string)
	String() string
	ToJSON() string
}

type Branch struct {
	parent *Branch
	key    string
	kids   []Node
}

func NewRoot() *Branch {
	return NewBranch("")
}

func NewBranch(key string) *Branch {
	return &Branch{
		key:  key,
		kids: []Node{},
	}
}

func (b *Branch) String() string {
	str := ""
	prevTyp := NoType
	lvl := b.Level()
	if lvl > 0 {
		str = strings.Repeat("=", lvl) + " " + b.key + "\n"
		prevTyp = CommentType // fake type to get a newline between branch and first kid
	}
	for _, kid := range b.kids {
		typ := kid.Type()
		if prevTyp == BranchType && typ != BranchType {
			str += "\n" + strings.Repeat("=", lvl+1) + "\n"
		}
		if prevTyp != NoType && !(typ == LeafType && prevTyp == LeafType) {
			str += "\n"
		}
		if typ == TextType && (prevTyp == TextType || prevTyp == LongLeafType) {
			str += "-\n\n"
		}
		str += kid.String()
		prevTyp = typ
	}
	return str
}

func (b *Branch) ToJSON() string {
	str := "{ \"" + b.key + "\": [ "
	for i, kid := range b.kids {
		str += kid.ToJSON()
		if i < len(b.kids)-1 {
			str += ","
		}
		str += " "
	}
	str += "] }"
	return str
}

func (b *Branch) Type() int {
	return BranchType
}

func (b *Branch) Level() int {
	i := 0
	for ; b.parent != nil; i++ {
		b = b.parent
	}
	return i
}

func (b *Branch) Key() string       { return b.key }
func (b *Branch) SetKey(key string) { b.key = key }

func (b *Branch) Parent() *Branch          { return b.parent }
func (b *Branch) SetParent(parent *Branch) { b.parent = parent }

func (b *Branch) NumKids() int { return len(b.kids) }
func (b *Branch) Kids() []Node { return b.kids }

func (b *Branch) Append(node Node) {
	b.kids = append(b.kids, node)
	node.SetParent(b)
}

func (b *Branch) First() Node {
	if len(b.kids) == 0 {
		return nil
	}
	return b.kids[0]
}

func (b *Branch) Insert(node Node, pos int) bool {
	if pos > len(b.kids) {
		return false
	}
	kids := append(b.kids[:pos], node)
	kids = append(kids, b.kids[pos:]...)
	b.kids = kids
	return true
}

type Leaf struct {
	typ    int
	key    string
	val    string
	parent *Branch
}

func NewLeaf(key, val string) *Leaf {
	return &Leaf{
		key: key,
		val: val,
		typ: LeafType,
	}
}

func NewLongLeaf(key, val string) *Leaf {
	leaf := NewLeaf(key, val)
	leaf.typ = LongLeafType
	return leaf
}

func NewComment(val string) *Leaf {
	leaf := NewLeaf("#", val)
	leaf.typ = CommentType
	return leaf
}

func (l *Leaf) String() string {
	str := ""
	switch l.typ {
	case LeafType:
		str += "-"
		if l.key != "" {
			str += " " + l.key
		}
		if l.val != "" {
			str += ": " + l.val
		}
		str += "\n"
	case LongLeafType:
		str = "- " + l.key + ":\n\n"
		str += l.val
	case TextType, CommentType:
		str += l.val
	}
	return str
}

func (l *Leaf) ToJSON() string {
	val := strings.Replace(l.val, "\n", "\\n", -1)
	val = strings.Replace(val, "\"", "\\\"", -1)
	return "{ \"" + l.key + "\": \"" + val + "\" }"
}

func (l *Leaf) Parent() *Branch          { return l.parent }
func (l *Leaf) SetParent(branch *Branch) { l.parent = branch }

func (l *Leaf) Type() int       { return l.typ }
func (l *Leaf) SetType(typ int) { l.typ = typ }

func (l *Leaf) Key() string       { return l.key }
func (l *Leaf) SetKey(key string) { l.key = key }

func (l *Leaf) Value() string          { return l.val }
func (l *Leaf) SetValue(val string)    { l.val = val }
func (l *Leaf) AppendValue(val string) { l.val += val }
<commit_msg>Added bufio-based output for speed.<commit_after>package dgrl

import (
	"bufio"
	"errors"
	"io"
	"strings"
)

type Node interface {
	Parent() *Branch
	SetParent(parent *Branch)
	Type() int
	Key() string
	SetKey(key string)
	String() string
	ToJSON() string

	write(w *bufio.Writer)
}

type Branch struct {
	parent *Branch
	key    string
	kids   []Node
}

func NewRoot() *Branch {
	return NewBranch("")
}

func NewBranch(key string) *Branch {
	return &Branch{
		key:  key,
		kids: []Node{},
	}
}

func (b *Branch) String() string {
	str := ""
	prevTyp := NoType
	lvl := b.Level()
	if lvl > 0 {
		str = strings.Repeat("=", lvl) + " " + b.key + "\n"
		prevTyp = CommentType // fake type to get a newline between branch and first kid
	}
	for _, kid := range b.kids {
		typ := kid.Type()
		if prevTyp == BranchType && typ != BranchType {
			str += "\n" + strings.Repeat("=", lvl+1) + "\n"
		}
		if prevTyp != NoType && !(typ == LeafType && prevTyp == LeafType) {
			str += "\n"
		}
		if typ == TextType && (prevTyp == TextType || prevTyp == LongLeafType) {
			str += "-\n\n"
		}
		str += kid.String()
		prevTyp = typ
	}
	return str
}

func (b *Branch) Write(w io.Writer) error {
	buf := bufio.NewWriter(w)
	if buf == nil {
		return errors.New("error creating bufio.Writer")
	}
	b.write(buf)
	buf.Flush()
	return nil
}

func (b *Branch) write(w *bufio.Writer) {
	prevTyp := NoType
	lvl := b.Level()
	if lvl > 0 {
		for i := 0; i < lvl; i++ {
			w.WriteByte('=')
		}
		w.WriteByte(' ')
		w.WriteString(b.key)
		prevTyp = CommentType // fake type to get a newline between branch and first kid
	}
	for _, kid := range b.kids {
		typ := kid.Type()
		if prevTyp == BranchType && typ != BranchType {
			w.WriteByte('\n')
			for i := 0; i < lvl+1; i++ {
				w.WriteByte('=')
			}
			w.WriteByte('\n')
		}
		if prevTyp != NoType && !(typ == LeafType && prevTyp == LeafType) {
			w.WriteByte('\n')
		}
		if typ == TextType && (prevTyp == TextType || prevTyp == LongLeafType) {
			w.WriteString("-\n\n")
		}
		kid.write(w)
		prevTyp = typ
	}
}

func (b *Branch) ToJSON() string {
	str := "{ \"" + b.key + "\": [ "
	for i, kid := range b.kids {
		str += kid.ToJSON()
		if i < len(b.kids)-1 {
			str += ","
		}
		str += " "
	}
	str += "] }"
	return str
}

func (b *Branch) Type() int {
	return BranchType
}

func (b *Branch) Level() int {
	i := 0
	for ; b.parent != nil; i++ {
		b = b.parent
	}
	return i
}

func (b *Branch) Key() string       { return b.key }
func (b *Branch) SetKey(key string) { b.key = key }

func (b *Branch) Parent() *Branch          { return b.parent }
func (b *Branch) SetParent(parent *Branch) { b.parent = parent }

func (b *Branch) NumKids() int { return len(b.kids) }
func (b *Branch) Kids() []Node { return b.kids }

func (b *Branch) Append(node Node) {
	b.kids = append(b.kids, node)
	node.SetParent(b)
}

func (b *Branch) First() Node {
	if len(b.kids) == 0 {
		return nil
	}
	return b.kids[0]
}

func (b *Branch) Insert(node Node, pos int) bool {
	if pos > len(b.kids) {
		return false
	}
	kids := append(b.kids[:pos], node)
	kids = append(kids, b.kids[pos:]...)
	b.kids = kids
	return true
}

type Leaf struct {
	typ    int
	key    string
	val    string
	parent *Branch
}

func NewLeaf(key, val string) *Leaf {
	return &Leaf{
		key: key,
		val: val,
		typ: LeafType,
	}
}

func NewLongLeaf(key, val string) *Leaf {
	leaf := NewLeaf(key, val)
	leaf.typ = LongLeafType
	return leaf
}

func NewComment(val string) *Leaf {
	leaf := NewLeaf("#", val)
	leaf.typ = CommentType
	return leaf
}

func (l *Leaf) String() string {
	str := ""
	switch l.typ {
	case LeafType:
		str += "-"
		if l.key != "" {
			str += " " + l.key
		}
		if l.val != "" {
			str += ": " + l.val
		}
		str += "\n"
	case LongLeafType:
		str = "- " + l.key + ":\n\n"
		str += l.val
	case TextType, CommentType:
		str += l.val
	}
	return str
}

func (l *Leaf) write(w *bufio.Writer) {
	switch l.typ {
	case LeafType:
		w.WriteByte('-')
		if l.key != "" {
			w.WriteByte(' ')
			w.WriteString(l.key)
		}
		if l.val != "" {
			w.WriteString(": ")
			w.WriteString(l.val)
		}
		w.WriteByte('\n')
	case LongLeafType:
		w.WriteString("- ")
		w.WriteString(l.key)
		w.WriteString(":\n\n")
		w.WriteString(l.val)
	case TextType, CommentType:
		w.WriteString(l.val)
	}
}

func (l *Leaf) ToJSON() string {
	val := strings.Replace(l.val, "\n", "\\n", -1)
	val = strings.Replace(val, "\"", "\\\"", -1)
	return "{ \"" + l.key + "\": \"" + val + "\" }"
}

func (l *Leaf) Parent() *Branch          { return l.parent }
func (l *Leaf) SetParent(branch *Branch) { l.parent = branch }

func (l *Leaf) Type() int       { return l.typ }
func (l *Leaf) SetType(typ int) { l.typ = typ }

func (l *Leaf) Key() string       { return l.key }
func (l *Leaf) SetKey(key string) { l.key = key }

func (l *Leaf) Value() string          { return l.val }
func (l *Leaf) SetValue(val string)    { l.val = val }
func (l *Leaf) AppendValue(val string) { l.val += val }
<|endoftext|>


<commit_before>package javaclass

import (
	"io"

	"github.com/MJKWoolnough/byteio"
)

const (
	ConstantNull               = 0
	ConstantUTF8               = 1
	ConstantInteger            = 3
	ConstantFloat              = 4
	ConstantLong               = 5
	ConstantDouble             = 6
	ConstantClass              = 7
	ConstantString             = 8
	ConstantFieldRef           = 9
	ConstantMethodRef          = 10
	ConstantInterfaceMethodRef = 11
	ConstantNameAndType        = 12
	ConstantMethodHandle       = 15
	ConstantMethodType         = 16
	ConstantInvokeDynamic      = 18
)

type CPInfo interface {
	Type() int
}

func readConstantPool(r io.Reader) ([]CPInfo, error) {
	br := byteio.BigEndianReader{r}
	constantPoolCount, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	constantPool := make([]CPInfo, 1, constantPoolCount)
	constantPool[0] = ConstantNullInfo{}
	for i := uint16(1); i < constantPoolCount; i++ {
		tag, _, err := br.ReadUint8()
		if err != nil {
			return nil, err
		}
		var (
			cpInfo  CPInfo
			addNull bool
		)
		switch tag {
		case ConstantUTF8:
			cpInfo, err = readConstantUTF8(r)
		case ConstantInteger:
			cpInfo, err = readConstantInteger(r)
		case ConstantFloat:
			cpInfo, err = readConstantFloat(r)
		case ConstantLong:
			cpInfo, err = readConstantLong(r)
		case ConstantDouble:
			cpInfo, err = readConstantDouble(r)
		case ConstantClass:
			cpInfo, err = readConstantClass(r)
		case ConstantString:
			cpInfo, err = readConstantString(r)
		case ConstantFieldRef:
			cpInfo, err = readConstantFieldRef(r)
		case ConstantMethodRef:
			cpInfo, err = readConstantMethodRef(r)
		case ConstantInterfaceMethodRef:
			cpInfo, err = readConstantInterfaceMethodRef(r)
		case ConstantNameAndType:
			cpInfo, err = readConstantNameAndType(r)
		case ConstantMethodHandle:
			cpInfo, err = readConstantMethodHandle(r)
		case ConstantMethodType:
			cpInfo, err = readConstantMethodType(r)
		case ConstantInvokeDynamic:
			cpInfo, err = readConstantInvokeDynamic(r)
		default:
			err = ErrUnknownConstantPoolTag{tag}
		}
		if err != nil {
			return nil, err
		}
		constantPool = append(constantPool, cpInfo)
		if addNull {
			constantPool = append(constantPool, ConstantNullInfo{})
		}
	}
	return constantPool, nil
}

type ConstantNullInfo struct{}

func (ConstantNullInfo) Type() int {
	return ConstantNull
}

type ConstantUTF8Info struct {
	String string
}

func readConstantUTF8(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	length, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	toString := make([]byte, length)
	_, err = br.Read(toString)
	if err != nil {
		return nil, err
	}
	return ConstantUTF8Info{string(toString)}, nil
}

func (ConstantUTF8Info) Type() int {
	return ConstantUTF8
}

type ConstantIntegerInfo struct {
	Integer uint32
}

func readConstantInteger(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint32()
	if err != nil {
		return nil, err
	}
	return ConstantIntegerInfo{n}, nil
}

func (ConstantIntegerInfo) Type() int {
	return ConstantInteger
}

type ConstantFloatInfo struct {
	Float float32
}

func readConstantFloat(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadFloat32()
	if err != nil {
		return nil, err
	}
	return ConstantFloatInfo{n}, nil
}

func (ConstantFloatInfo) Type() int {
	return ConstantFloat
}

type ConstantLongInfo struct {
	Long uint64
}

func readConstantLong(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint64()
	if err != nil {
		return nil, err
	}
	return ConstantLongInfo{n}, nil
}

func (ConstantLongInfo) Type() int {
	return ConstantLong
}

type ConstantDoubleInfo struct {
	Double float64
}

func readConstantDouble(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadFloat64()
	if err != nil {
		return nil, err
	}
	return ConstantDoubleInfo{n}, nil
}

func (ConstantDoubleInfo) Type() int {
	return ConstantDouble
}

type ConstantClassInfo struct {
	NameIndex uint16
}

func readConstantClass(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantClassInfo{n}, nil
}

func (ConstantClassInfo) Type() int {
	return ConstantClass
}

type ConstantStringInfo struct {
	StringIndex uint16
}

func readConstantString(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantStringInfo{n}, nil
}

func (ConstantStringInfo) Type() int {
	return ConstantString
}

type ConstantFieldRefInfo struct {
	ClassIndex, NameAndTypeIndex uint16
}

func readConstantFieldRef(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantFieldRefInfo{i, n}, nil
}

func (ConstantFieldRefInfo) Type() int {
	return ConstantFieldRef
}

type ConstantMethodRefInfo struct {
	ClassIndex, NameAndTypeIndex uint16
}

func readConstantMethodRef(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantMethodRefInfo{i, n}, nil
}

func (ConstantMethodRefInfo) Type() int {
	return ConstantMethodRef
}

type ConstantInterfaceMethodRefInfo struct {
	ClassIndex, NameAndTypeIndex uint16
}

func readConstantInterfaceMethodRef(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantInterfaceMethodRefInfo{i, n}, nil
}

func (ConstantInterfaceMethodRefInfo) Type() int {
	return ConstantInterfaceMethodRef
}

type ConstantNameAndTypeInfo struct {
	NameIndex, DescriptorIndex uint16
}

func readConstantNameAndType(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	d, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantNameAndTypeInfo{n, d}, nil
}

func (ConstantNameAndTypeInfo) Type() int {
	return ConstantNameAndType
}

type ConstantMethodHandleInfo struct {
	ReferenceKind  uint8
	ReferenceIndex uint16
}

func readConstantMethodHandle(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	k, _, err := br.ReadUint8()
	if err != nil {
		return nil, err
	}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantMethodHandleInfo{k, i}, nil
}

func (ConstantMethodHandleInfo) Type() int {
	return ConstantMethodHandle
}

type ConstantMethodTypeInfo struct {
	DescriptorIndex uint16
}

func readConstantMethodType(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantMethodTypeInfo{i}, nil
}

func (ConstantMethodTypeInfo) Type() int {
	return ConstantMethodType
}

type ConstantInvokeDynamicInfo struct {
	BootstrapMethodAttrIndex, NameAndTypeIndex uint16
}

func readConstantInvokeDynamic(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	b, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantInvokeDynamicInfo{b, i}, nil
}

func (ConstantInvokeDynamicInfo) Type() int {
	return ConstantInvokeDynamic
}

// Error types

type ErrUnknownConstantPoolTag struct {
	Tag uint8
}

func (ErrUnknownConstantPoolTag) Error() string {
	return "unknown constant pool tag"
}
<commit_msg>Added missing position incrementation for double and long<commit_after>package javaclass

import (
	"io"

	"github.com/MJKWoolnough/byteio"
)

const (
	ConstantNull               = 0
	ConstantUTF8               = 1
	ConstantInteger            = 3
	ConstantFloat              = 4
	ConstantLong               = 5
	ConstantDouble             = 6
	ConstantClass              = 7
	ConstantString             = 8
	ConstantFieldRef           = 9
	ConstantMethodRef          = 10
	ConstantInterfaceMethodRef = 11
	ConstantNameAndType        = 12
	ConstantMethodHandle       = 15
	ConstantMethodType         = 16
	ConstantInvokeDynamic      = 18
)

type CPInfo interface {
	Type() int
}

func readConstantPool(r io.Reader) ([]CPInfo, error) {
	br := byteio.BigEndianReader{r}
	constantPoolCount, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	constantPool := make([]CPInfo, 1, constantPoolCount)
	constantPool[0] = ConstantNullInfo{}
	for i := uint16(1); i < constantPoolCount; i++ {
		tag, _, err := br.ReadUint8()
		if err != nil {
			return nil, err
		}
		var (
			cpInfo  CPInfo
			addNull bool
		)
		switch tag {
		case ConstantUTF8:
			cpInfo, err = readConstantUTF8(r)
		case ConstantInteger:
			cpInfo, err = readConstantInteger(r)
		case ConstantFloat:
			cpInfo, err = readConstantFloat(r)
		case ConstantLong:
			cpInfo, err = readConstantLong(r)
			addNull = true
		case ConstantDouble:
			cpInfo, err = readConstantDouble(r)
			addNull = true
		case ConstantClass:
			cpInfo, err = readConstantClass(r)
		case ConstantString:
			cpInfo, err = readConstantString(r)
		case ConstantFieldRef:
			cpInfo, err = readConstantFieldRef(r)
		case ConstantMethodRef:
			cpInfo, err = readConstantMethodRef(r)
		case ConstantInterfaceMethodRef:
			cpInfo, err = readConstantInterfaceMethodRef(r)
		case ConstantNameAndType:
			cpInfo, err = readConstantNameAndType(r)
		case ConstantMethodHandle:
			cpInfo, err = readConstantMethodHandle(r)
		case ConstantMethodType:
			cpInfo, err = readConstantMethodType(r)
		case ConstantInvokeDynamic:
			cpInfo, err = readConstantInvokeDynamic(r)
		default:
			err = ErrUnknownConstantPoolTag{tag}
		}
		if err != nil {
			return nil, err
		}
		constantPool = append(constantPool, cpInfo)
		if addNull {
			constantPool = append(constantPool, ConstantNullInfo{})
			i++
		}
	}
	return constantPool, nil
}

type ConstantNullInfo struct{}

func (ConstantNullInfo) Type() int {
	return ConstantNull
}

type ConstantUTF8Info struct {
	String string
}

func readConstantUTF8(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	length, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	toString := make([]byte, length)
	_, err = br.Read(toString)
	if err != nil {
		return nil, err
	}
	return ConstantUTF8Info{string(toString)}, nil
}

func (ConstantUTF8Info) Type() int {
	return ConstantUTF8
}

type ConstantIntegerInfo struct {
	Integer uint32
}

func readConstantInteger(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint32()
	if err != nil {
		return nil, err
	}
	return ConstantIntegerInfo{n}, nil
}

func (ConstantIntegerInfo) Type() int {
	return ConstantInteger
}

type ConstantFloatInfo struct {
	Float float32
}

func readConstantFloat(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadFloat32()
	if err != nil {
		return nil, err
	}
	return ConstantFloatInfo{n}, nil
}

func (ConstantFloatInfo) Type() int {
	return ConstantFloat
}

type ConstantLongInfo struct {
	Long uint64
}

func readConstantLong(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint64()
	if err != nil {
		return nil, err
	}
	return ConstantLongInfo{n}, nil
}

func (ConstantLongInfo) Type() int {
	return ConstantLong
}

type ConstantDoubleInfo struct {
	Double float64
}

func readConstantDouble(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadFloat64()
	if err != nil {
		return nil, err
	}
	return ConstantDoubleInfo{n}, nil
}

func (ConstantDoubleInfo) Type() int {
	return ConstantDouble
}

type ConstantClassInfo struct {
	NameIndex uint16
}

func readConstantClass(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantClassInfo{n}, nil
}

func (ConstantClassInfo) Type() int {
	return ConstantClass
}

type ConstantStringInfo struct {
	StringIndex uint16
}

func readConstantString(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantStringInfo{n}, nil
}

func (ConstantStringInfo) Type() int {
	return ConstantString
}

type ConstantFieldRefInfo struct {
	ClassIndex, NameAndTypeIndex uint16
}

func readConstantFieldRef(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantFieldRefInfo{i, n}, nil
}

func (ConstantFieldRefInfo) Type() int {
	return ConstantFieldRef
}

type ConstantMethodRefInfo struct {
	ClassIndex, NameAndTypeIndex uint16
}

func readConstantMethodRef(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantMethodRefInfo{i, n}, nil
}

func (ConstantMethodRefInfo) Type() int {
	return ConstantMethodRef
}

type ConstantInterfaceMethodRefInfo struct {
	ClassIndex, NameAndTypeIndex uint16
}

func readConstantInterfaceMethodRef(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantInterfaceMethodRefInfo{i, n}, nil
}

func (ConstantInterfaceMethodRefInfo) Type() int {
	return ConstantInterfaceMethodRef
}

type ConstantNameAndTypeInfo struct {
	NameIndex, DescriptorIndex uint16
}

func readConstantNameAndType(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	n, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	d, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantNameAndTypeInfo{n, d}, nil
}

func (ConstantNameAndTypeInfo) Type() int {
	return ConstantNameAndType
}

type ConstantMethodHandleInfo struct {
	ReferenceKind  uint8
	ReferenceIndex uint16
}

func readConstantMethodHandle(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	k, _, err := br.ReadUint8()
	if err != nil {
		return nil, err
	}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantMethodHandleInfo{k, i}, nil
}

func (ConstantMethodHandleInfo) Type() int {
	return ConstantMethodHandle
}

type ConstantMethodTypeInfo struct {
	DescriptorIndex uint16
}

func readConstantMethodType(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantMethodTypeInfo{i}, nil
}

func (ConstantMethodTypeInfo) Type() int {
	return ConstantMethodType
}

type ConstantInvokeDynamicInfo struct {
	BootstrapMethodAttrIndex, NameAndTypeIndex uint16
}

func readConstantInvokeDynamic(r io.Reader) (CPInfo, error) {
	br := byteio.BigEndianReader{r}
	b, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	i, _, err := br.ReadUint16()
	if err != nil {
		return nil, err
	}
	return ConstantInvokeDynamicInfo{b, i}, nil
}

func (ConstantInvokeDynamicInfo) Type() int {
	return ConstantInvokeDynamic
}

// Error types

type ErrUnknownConstantPoolTag struct {
	Tag uint8
}

func (ErrUnknownConstantPoolTag) Error() string {
	return "unknown constant pool tag"
}
<|endoftext|>


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

package blob_test

import (
	"errors"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	. "github.com/jacobsa/oglematchers"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestChecking(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type checkingStoreTest struct {
	wrapped mock_blob.MockStore
	store   blob.Store
}

func (t *checkingStoreTest) SetUp(i *TestInfo) {
	t.wrapped = mock_blob.NewMockStore(i.MockController, "wrapped")
	t.store = blob.NewCheckingStore(t.wrapped)
}

////////////////////////////////////////////////////////////////////////
// Store
////////////////////////////////////////////////////////////////////////

type CheckingStore_StoreTest struct {
	checkingStoreTest
}

func init() { RegisterTestSuite(&CheckingStore_StoreTest{}) }

func (t *CheckingStore_StoreTest) CallsWrapped() {
	b := []byte{0xde, 0xad}

	// Wrapped
	ExpectCall(t.wrapped, "Store")(DeepEquals(b)).
		WillOnce(oglemock.Return(blob.Score{}, errors.New("")))

	// Call
	t.store.Store(b)
}

func (t *CheckingStore_StoreTest) WrappedReturnsError() {
	b := []byte{}

	// Wrapped
	ExpectCall(t.wrapped, "Store")(Any()).
		WillOnce(oglemock.Return(blob.Score{}, errors.New("taco")))

	// Call
	_, err := t.store.Store(b)

	ExpectThat(err, Error(Equals("taco")))
}

func (t *CheckingStore_StoreTest) WrappedReturnsIncorrectScore() {
	b := []byte{0xde, 0xad}
	correctScore := blob.ComputeScore(b)
	incorrectScore := blob.ComputeScore([]byte{0xbe, 0xef})

	// Wrapped
	ExpectCall(t.wrapped, "Store")(Any()).
		WillOnce(oglemock.Return(incorrectScore, nil))

	// Call
	_, err := t.store.Store(b)

	ExpectThat(err, Error(HasSubstr("Incorrect")))
	ExpectThat(err, Error(HasSubstr("score")))
	ExpectThat(err, Error(HasSubstr(correctScore.Hex())))
	ExpectThat(err, Error(HasSubstr(incorrectScore.Hex())))
}

func (t *CheckingStore_StoreTest) WrappedReturnsCorrectScore() {
	b := []byte{0xde, 0xad}
	correctScore := blob.ComputeScore(b)

	// Wrapped
	ExpectCall(t.wrapped, "Store")(Any()).
		WillOnce(oglemock.Return(correctScore, nil))

	// Call
	score, err := t.store.Store(b)

	AssertEq(nil, err)
	ExpectThat(score, DeepEquals(correctScore))
}

////////////////////////////////////////////////////////////////////////
// Load
////////////////////////////////////////////////////////////////////////

type CheckingStore_LoadTest struct {
	checkingStoreTest
}

func init() { RegisterTestSuite(&CheckingStore_LoadTest{}) }

func (t *CheckingStore_LoadTest) CallsWrapped() {
	score := blob.ComputeScore([]byte{0xde, 0xad})

	// Wrapped
	ExpectCall(t.wrapped, "Load")(DeepEquals(score)).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.store.Load(score)
}

func (t *CheckingStore_LoadTest) WrappedReturnsError() {
	score := blob.ComputeScore([]byte{})

	// Wrapped
	ExpectCall(t.wrapped, "Load")(Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	_, err := t.store.Load(score)

	ExpectThat(err, Error(Equals("taco")))
}

func (t *CheckingStore_LoadTest) WrappedReturnsIncorrectData() {
	correctData := []byte{0xde, 0xad}
	incorrectData := []byte{0xde, 0xad, 0xbe, 0xef}
	score := blob.ComputeScore(correctData)

	// Wrapped
	ExpectCall(t.wrapped, "Load")(Any()).
		WillOnce(oglemock.Return(incorrectData, nil))

	// Call
	_, err := t.store.Load(score)

	ExpectThat(err, Error(HasSubstr("Incorrect")))
	ExpectThat(err, Error(HasSubstr("data")))
	ExpectThat(err, Error(HasSubstr(score.Hex())))
}

func (t *CheckingStore_LoadTest) WrappedReturnsCorrectData() {
	correctData := []byte{0xde, 0xad}
	score := blob.ComputeScore(correctData)

	// Wrapped
	ExpectCall(t.wrapped, "Load")(Any()).
		WillOnce(oglemock.Return(correctData, nil))

	// Call
	data, err := t.store.Load(score)

	AssertEq(nil, err)
	ExpectThat(data, DeepEquals(correctData))
}
<commit_msg>checking_test.go<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package blob_test

import (
	"errors"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	. "github.com/jacobsa/oglematchers"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
)

func TestChecking(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type checkingStoreTest struct {
	ctx     context.Context
	wrapped mock_blob.MockStore
	store   blob.Store
}

func (t *checkingStoreTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
	t.wrapped = mock_blob.NewMockStore(ti.MockController, "wrapped")
	t.store = blob.NewCheckingStore(t.wrapped)
}

////////////////////////////////////////////////////////////////////////
// Store
////////////////////////////////////////////////////////////////////////

type CheckingStore_StoreTest struct {
	checkingStoreTest
}

func init() { RegisterTestSuite(&CheckingStore_StoreTest{}) }

func (t *CheckingStore_StoreTest) CallsWrapped() {
	b := []byte{0xde, 0xad}

	// Wrapped
	ExpectCall(t.wrapped, "Store")(Any(), DeepEquals(b)).
		WillOnce(oglemock.Return(blob.Score{}, errors.New("")))

	// Call
	t.store.Store(t.ctx, b)
}

func (t *CheckingStore_StoreTest) WrappedReturnsError() {
	b := []byte{}

	// Wrapped
	ExpectCall(t.wrapped, "Store")(Any(), Any()).
		WillOnce(oglemock.Return(blob.Score{}, errors.New("taco")))

	// Call
	_, err := t.store.Store(t.ctx, b)

	ExpectThat(err, Error(Equals("taco")))
}

func (t *CheckingStore_StoreTest) WrappedReturnsIncorrectScore() {
	b := []byte{0xde, 0xad}
	correctScore := blob.ComputeScore(b)
	incorrectScore := blob.ComputeScore([]byte{0xbe, 0xef})

	// Wrapped
	ExpectCall(t.wrapped, "Store")(Any(), Any()).
		WillOnce(oglemock.Return(incorrectScore, nil))

	// Call
	_, err := t.store.Store(t.ctx, b)

	ExpectThat(err, Error(HasSubstr("Incorrect")))
	ExpectThat(err, Error(HasSubstr("score")))
	ExpectThat(err, Error(HasSubstr(correctScore.Hex())))
	ExpectThat(err, Error(HasSubstr(incorrectScore.Hex())))
}

func (t *CheckingStore_StoreTest) WrappedReturnsCorrectScore() {
	b := []byte{0xde, 0xad}
	correctScore := blob.ComputeScore(b)

	// Wrapped
	ExpectCall(t.wrapped, "Store")(Any(), Any()).
		WillOnce(oglemock.Return(correctScore, nil))

	// Call
	score, err := t.store.Store(t.ctx, b)

	AssertEq(nil, err)
	ExpectThat(score, DeepEquals(correctScore))
}

////////////////////////////////////////////////////////////////////////
// Load
////////////////////////////////////////////////////////////////////////

type CheckingStore_LoadTest struct {
	checkingStoreTest
}

func init() { RegisterTestSuite(&CheckingStore_LoadTest{}) }

func (t *CheckingStore_LoadTest) CallsWrapped() {
	score := blob.ComputeScore([]byte{0xde, 0xad})

	// Wrapped
	ExpectCall(t.wrapped, "Load")(Any(), DeepEquals(score)).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.store.Load(t.ctx, score)
}

func (t *CheckingStore_LoadTest) WrappedReturnsError() {
	score := blob.ComputeScore([]byte{})

	// Wrapped
	ExpectCall(t.wrapped, "Load")(Any(), Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	_, err := t.store.Load(t.ctx, score)

	ExpectThat(err, Error(Equals("taco")))
}

func (t *CheckingStore_LoadTest) WrappedReturnsIncorrectData() {
	correctData := []byte{0xde, 0xad}
	incorrectData := []byte{0xde, 0xad, 0xbe, 0xef}
	score := blob.ComputeScore(correctData)

	// Wrapped
	ExpectCall(t.wrapped, "Load")(Any(), Any()).
		WillOnce(oglemock.Return(incorrectData, nil))

	// Call
	_, err := t.store.Load(t.ctx, score)

	ExpectThat(err, Error(HasSubstr("Incorrect")))
	ExpectThat(err, Error(HasSubstr("data")))
	ExpectThat(err, Error(HasSubstr(score.Hex())))
}

func (t *CheckingStore_LoadTest) WrappedReturnsCorrectData() {
	correctData := []byte{0xde, 0xad}
	score := blob.ComputeScore(correctData)

	// Wrapped
	ExpectCall(t.wrapped, "Load")(Any(), Any()).
		WillOnce(oglemock.Return(correctData, nil))

	// Call
	data, err := t.store.Load(t.ctx, score)

	AssertEq(nil, err)
	ExpectThat(data, DeepEquals(correctData))
}
<|endoftext|>


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

package rand_test

import (
	"crypto/rand"
	"testing"
)

// http://golang.org/issue/6849.
func TestPrimeSmall(t *testing.T) {
	for n := 2; n < 10; n++ {
		p, err := rand.Prime(rand.Reader, n)
		if err != nil {
			t.Fatalf("Can't generate %d-bit prime: %v", n, err)
		}
		if p.BitLen() != n {
			t.Fatalf("%v is not %d-bit", p, n)
		}
		if !p.ProbablyPrime(32) {
			t.Fatalf("%v is not prime", p)
		}
	}
}
<commit_msg>crypto/rand: add tests for Int, Prime<commit_after>// Copyright 2013 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rand_test

import (
	"crypto/rand"
	"math/big"
	"testing"
)

// http://golang.org/issue/6849.
func TestPrimeSmall(t *testing.T) {
	for n := 2; n < 10; n++ {
		p, err := rand.Prime(rand.Reader, n)
		if err != nil {
			t.Fatalf("Can't generate %d-bit prime: %v", n, err)
		}
		if p.BitLen() != n {
			t.Fatalf("%v is not %d-bit", p, n)
		}
		if !p.ProbablyPrime(32) {
			t.Fatalf("%v is not prime", p)
		}
	}
}

// Test that passing bits < 2 causes Prime to return nil, error
func TestPrimeBitsLt2(t *testing.T) {
	if p, err := rand.Prime(rand.Reader, 1); p != nil || err == nil {
		t.Errorf("Prime should return nil, error when called with bits < 2")
	}
}

func TestInt(t *testing.T) {
	// start at 128 so the case of (max.BitLen() % 8) == 0 is covered
	for n := 128; n < 140; n++ {
		b := new(big.Int).SetInt64(int64(n))
		if i, err := rand.Int(rand.Reader, b); err != nil {
			t.Fatalf("Can't generate random value: %v, %v", i, err)
		}
	}
}

func testIntPanics(t *testing.T, b *big.Int) {
	defer func() {
		if err := recover(); err == nil {
			t.Errorf("Int should panic when called with max <= 0: %v", b)
		}
	}()
	rand.Int(rand.Reader, b)
}

// Test that passing a new big.Int as max causes Int to panic
func TestIntEmptyMaxPanics(t *testing.T) {
	b := new(big.Int)
	testIntPanics(t, b)
}

// Test that passing a negative value as max causes Int to panic
func TestIntNegativeMaxPanics(t *testing.T) {
	b := new(big.Int).SetInt64(int64(-1))
	testIntPanics(t, b)
}
<|endoftext|>


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

// Package iotest implements Readers and Writers useful only for testing.
package iotest

import (
	"errors"
	"io"
)

// OneByteReader returns a Reader that implements
// each non-empty Read by reading one byte from r.
func OneByteReader(r io.Reader) io.Reader { return &oneByteReader{r} }

type oneByteReader struct {
	r io.Reader
}

func (r *oneByteReader) Read(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	return r.r.Read(p[0:1])
}

// HalfReader returns a Reader that implements Read
// by reading half as many requested bytes from r.
func HalfReader(r io.Reader) io.Reader { return &halfReader{r} }

type halfReader struct {
	r io.Reader
}

func (r *halfReader) Read(p []byte) (int, error) {
	return r.r.Read(p[0 : (len(p)+1)/2])
}

// DataErrReader returns a Reader that returns the final
// error with the last data read, instead of by itself with
// zero bytes of data.
func DataErrReader(r io.Reader) io.Reader { return &dataErrReader{r, nil, make([]byte, 1024)} }

type dataErrReader struct {
	r      io.Reader
	unread []byte
	data   []byte
}

func (r *dataErrReader) Read(p []byte) (n int, err error) {
	// loop because first call needs two reads:
	// one to get data and a second to look for an error.
	for {
		if len(r.unread) == 0 {
			n1, err1 := r.r.Read(r.data)
			r.unread = r.data[0:n1]
			err = err1
		}
		if n > 0 || err != nil {
			break
		}
		n = copy(p, r.unread)
		r.unread = r.unread[n:]
	}
	return
}

var ErrTimeout = errors.New("timeout")

// TimeoutReader returns ErrTimeout on the second read
// with no data.  Subsequent calls to read succeed.
func TimeoutReader(r io.Reader) io.Reader { return &timeoutReader{r, 0} }

type timeoutReader struct {
	r     io.Reader
	count int
}

func (r *timeoutReader) Read(p []byte) (int, error) {
	r.count++
	if r.count == 2 {
		return 0, ErrTimeout
	}
	return r.r.Read(p)
}
<commit_msg>testing/iotest: adjust adverbial imprecision in package comment<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package iotest implements Readers and Writers useful mainly for testing.
package iotest

import (
	"errors"
	"io"
)

// OneByteReader returns a Reader that implements
// each non-empty Read by reading one byte from r.
func OneByteReader(r io.Reader) io.Reader { return &oneByteReader{r} }

type oneByteReader struct {
	r io.Reader
}

func (r *oneByteReader) Read(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	return r.r.Read(p[0:1])
}

// HalfReader returns a Reader that implements Read
// by reading half as many requested bytes from r.
func HalfReader(r io.Reader) io.Reader { return &halfReader{r} }

type halfReader struct {
	r io.Reader
}

func (r *halfReader) Read(p []byte) (int, error) {
	return r.r.Read(p[0 : (len(p)+1)/2])
}

// DataErrReader returns a Reader that returns the final
// error with the last data read, instead of by itself with
// zero bytes of data.
func DataErrReader(r io.Reader) io.Reader { return &dataErrReader{r, nil, make([]byte, 1024)} }

type dataErrReader struct {
	r      io.Reader
	unread []byte
	data   []byte
}

func (r *dataErrReader) Read(p []byte) (n int, err error) {
	// loop because first call needs two reads:
	// one to get data and a second to look for an error.
	for {
		if len(r.unread) == 0 {
			n1, err1 := r.r.Read(r.data)
			r.unread = r.data[0:n1]
			err = err1
		}
		if n > 0 || err != nil {
			break
		}
		n = copy(p, r.unread)
		r.unread = r.unread[n:]
	}
	return
}

var ErrTimeout = errors.New("timeout")

// TimeoutReader returns ErrTimeout on the second read
// with no data.  Subsequent calls to read succeed.
func TimeoutReader(r io.Reader) io.Reader { return &timeoutReader{r, 0} }

type timeoutReader struct {
	r     io.Reader
	count int
}

func (r *timeoutReader) Read(p []byte) (int, error) {
	r.count++
	if r.count == 2 {
		return 0, ErrTimeout
	}
	return r.r.Read(p)
}
<|endoftext|>