<commit_before>// Copyright 2013-2014 Adam Presley. All rights reserved
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package storage

import (
	"database/sql"
	"fmt"
	"log"
	"strings"

	"github.com/adampresley/golangdb"
	"github.com/mailslurper/libmailslurper/model/attachment"
	"github.com/mailslurper/libmailslurper/model/mailitem"
	"github.com/mailslurper/libmailslurper/sanitization"
)

/*
Creates a global connection handle in a map named "lib".
*/
func ConnectToStorage(connectionInfo *golangdb.DatabaseConnection) error {
	var err error

	log.Println("libmailslurper: IFNO - Connecting to database", connectionInfo.Database)

	err = connectionInfo.Connect("lib")
	if err != nil {
		return err
	}

	switch connectionInfo.Engine {
	case golangdb.SQLITE:
		CreateSqlliteDatabase()
	}

	return nil
}

/*
Disconnects from the database storage
*/
func DisconnectFromStorage() {
	golangdb.Db["lib"].Close()
}

/*
Returns an attachment by ID
*/
func GetAttachment(mailId, attachmentId string) (attachment.Attachment, error) {
	result := attachment.Attachment{}

	rows, err := golangdb.Db["lib"].Query(`
		SELECT
			  fileName TEXT
			, contentType TEXT
			, content TEXT
		FROM attachment
		WHERE
			id=?
			AND mailItemId=?
	`, attachmentId, mailId)

	if err != nil {
		return result, fmt.Errorf("Error running query to get attachment")
	}

	defer rows.Close()
	rows.Next()

	var fileName string
	var contentType string
	var content string

	rows.Scan(&fileName, &contentType, &content)

	result.Headers = &attachment.AttachmentHeader{
		FileName:    fileName,
		ContentType: contentType,
	}

	result.Contents = content
	return result, nil
}

func getMailQuery(whereClause string) string {
	sql := `
		SELECT
			  mailitem.id AS mailItemId
			, mailitem.dateSent
			, mailitem.fromAddress
			, mailitem.toAddressList
			, mailitem.subject
			, mailitem.xmailer
			, mailitem.body
			, mailitem.contentType
			, mailitem.boundary

		FROM mailitem

		WHERE 1=1 `

	sql = sql + whereClause
	sql = sql + ` ORDER BY mailitem.dateSent DESC `

	return sql
}

/*
Returns a single mail item by ID.
*/
func GetMail(id string) (mailitem.MailItem, error) {
	result := mailitem.MailItem{}
	xssService := sanitization.NewXSSService()

	sql := getMailQuery(" AND mailitem.id=? ")
	rows, err := golangdb.Db["lib"].Query(sql, id)

	if err != nil {
		return result, fmt.Errorf("Error running query to get mail items: %s", err)
	}

	rows.Next()

	var mailItemId string
	var dateSent string
	var fromAddress string
	var toAddressList string
	var subject string
	var xmailer string
	var body string
	var contentType string
	var boundary string

	rows.Scan(&mailItemId, &dateSent, &fromAddress, &toAddressList, &subject, &xmailer, &body, &contentType, &boundary)

	result = mailitem.MailItem{
		Id:          mailItemId,
		DateSent:    dateSent,
		FromAddress: fromAddress,
		ToAddresses: strings.Split(toAddressList, "; "),
		Subject:     xssService.SanitizeString(subject),
		XMailer:     xssService.SanitizeString(xmailer),
		Body:        xssService.SanitizeString(body),
		ContentType: contentType,
		Boundary:    boundary,
	}

	/*
	 * Get attachments
	 */
	sql = `
		SELECT
			  attachment.id AS attachmentId
			, attachment.fileName
			, attachment.contentType

		FROM attachment
		WHERE
			attachment.mailItemId=?`

	attachmentRows, err := golangdb.Db["lib"].Query(sql, mailItemId)
	if err != nil {
		return result, err
	}

	attachments := make([]*attachment.Attachment, 0)

	for attachmentRows.Next() {
		var attachmentId string
		var fileName string
		var contentType string

		attachmentRows.Scan(&attachmentId, &fileName, &contentType)

		newAttachment := &attachment.Attachment{
			Id:     attachmentId,
			MailId: mailItemId,
			Headers: &attachment.AttachmentHeader{
				FileName:    xssService.SanitizeString(fileName),
				ContentType: contentType,
			},
		}

		attachments = append(attachments, newAttachment)
	}

	attachmentRows.Close()

	result.Attachments = attachments

	rows.Close()
	return result, nil
}

/*
Retrieves all stored mail items as an array of MailItem items. Only
returns rows starting at offset and gets up to length records. NOTE:
This code stinks. It gets ALL rows, then returns a slice and the total number
of mail items. Ick!
*/
func GetMailCollection(offset, length int) ([]mailitem.MailItem, int, error) {
	result := make([]mailitem.MailItem, 0)
	xssService := sanitization.NewXSSService()

	sql := getMailQuery("")
	rows, err := golangdb.Db["lib"].Query(sql)

	if err != nil {
		return result, 0, fmt.Errorf("Error running query to get mail items: %s", err)
	}

	/*
	 * Loop over our records and grab attachments on the way.
	 */
	for rows.Next() {
		var mailItemId string
		var dateSent string
		var fromAddress string
		var toAddressList string
		var subject string
		var xmailer string
		var body string
		var contentType string
		var boundary string

		rows.Scan(&mailItemId, &dateSent, &fromAddress, &toAddressList, &subject, &xmailer, &body, &contentType, &boundary)

		newItem := mailitem.MailItem{
			Id:          mailItemId,
			DateSent:    dateSent,
			FromAddress: fromAddress,
			ToAddresses: strings.Split(toAddressList, "; "),
			Subject:     xssService.SanitizeString(subject),
			XMailer:     xssService.SanitizeString(xmailer),
			Body:        xssService.SanitizeString(body),
			ContentType: contentType,
			Boundary:    boundary,
		}

		/*
		 * Get attachments
		 */
		sql = `
			SELECT
				  attachment.id AS attachmentId
				, attachment.fileName
			FROM attachment
			WHERE
				attachment.mailItemId=?`

		attachmentRows, err := golangdb.Db["lib"].Query(sql, mailItemId)
		if err != nil {
			return result, 0, err
		}

		attachments := make([]*attachment.Attachment, 0)

		for attachmentRows.Next() {
			var attachmentId string
			var fileName string

			attachmentRows.Scan(&attachmentId, &fileName)

			newAttachment := &attachment.Attachment{
				Id:      attachmentId,
				MailId:  mailItemId,
				Headers: &attachment.AttachmentHeader{FileName: xssService.SanitizeString(fileName)},
			}

			attachments = append(attachments, newAttachment)
		}

		attachmentRows.Close()

		newItem.Attachments = attachments
		result = append(result, newItem)
	}

	rows.Close()

	totalRecords := len(result)
	return result, totalRecords, nil
}

/*
GetMailCount returns a count of mail items in the storage system.
*/
func GetMailCount() (int, error) {
	var mailItemCount int
	var err error

	err = golangdb.Db["lib"].QueryRow(`SELECT COUNT(id) AS mailItemCount FROM mailitem`).Scan(&mailItemCount)
	if err != nil {
		return 0, fmt.Errorf("Error running query to get mail item count: %s", err)
	}

	return mailItemCount, nil
}

func storeAttachments(mailItemId string, transaction *sql.Tx, attachments []*attachment.Attachment) error {
	for _, a := range attachments {
		attachmentId, err := mailitem.GenerateId()
		if err != nil {
			return fmt.Errorf("Error generating ID for attachment: %s", err)
		}

		statement, err := transaction.Prepare(`
			INSERT INTO attachment (
				  id
				, mailItemId
				, fileName
				, contentType
				, content
			) VALUES (
				  ?
				, ?
				, ?
				, ?
				, ?
			)
		`)

		if err != nil {
			return fmt.Errorf("Error preparing insert attachment statement: %s", err)
		}

		_, err = statement.Exec(
			attachmentId,
			mailItemId,
			a.Headers.FileName,
			a.Headers.ContentType,
			a.Contents,
		)

		if err != nil {
			return fmt.Errorf("Error executing insert attachment in StoreMail: %s", err)
		}

		statement.Close()
		a.Id = attachmentId
	}

	return nil
}

func StoreMail(mailItem *mailitem.MailItem) (string, error) {
	/*
	 * Create a transaction and insert the new mail item
	 */
	transaction, err := golangdb.Db["lib"].Begin()
	if err != nil {
		return "", fmt.Errorf("Error starting transaction in StoreMail: %s", err)
	}

	/*
	 * Insert the mail item
	 */
	statement, err := transaction.Prepare(`
			INSERT INTO mailitem (
				  id
				, dateSent
				, fromAddress
				, toAddressList
				, subject
				, xmailer
				, body
				, contentType
				, boundary
			) VALUES (
				  ?
				, ?
				, ?
				, ?
				, ?
				, ?
				, ?
				, ?
				, ?
			)
		`)

	if err != nil {
		return "", fmt.Errorf("Error preparing insert statement for mail item in StoreMail: %s", err)
	}

	_, err = statement.Exec(
		mailItem.Id,
		mailItem.DateSent,
		mailItem.FromAddress,
		strings.Join(mailItem.ToAddresses, "; "),
		mailItem.Subject,
		mailItem.XMailer,
		mailItem.Body,
		mailItem.ContentType,
		mailItem.Boundary,
	)

	if err != nil {
		transaction.Rollback()
		return "", fmt.Errorf("Error executing insert for mail item in StoreMail: %s", err)
	}

	statement.Close()

	/*
	 * Insert attachments
	 */
	if err = storeAttachments(mailItem.Id, transaction, mailItem.Attachments); err != nil {
		transaction.Rollback()
		return "", fmt.Errorf("Unable to insert attachments in StoreMail: %s", err)
	}

	transaction.Commit()
	log.Printf("New mail item written to database.\n\n")

	return mailItem.Id, nil
}
<commit_msg>Added method to delete mails<commit_after>// Copyright 2013-2014 Adam Presley. All rights reserved
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package storage

import (
	"database/sql"
	"fmt"
	"log"
	"strings"

	"github.com/adampresley/golangdb"
	"github.com/mailslurper/libmailslurper/model/attachment"
	"github.com/mailslurper/libmailslurper/model/mailitem"
	"github.com/mailslurper/libmailslurper/sanitization"
)

/*
Creates a global connection handle in a map named "lib".
*/
func ConnectToStorage(connectionInfo *golangdb.DatabaseConnection) error {
	var err error

	log.Println("libmailslurper: IFNO - Connecting to database", connectionInfo.Database)

	err = connectionInfo.Connect("lib")
	if err != nil {
		return err
	}

	switch connectionInfo.Engine {
	case golangdb.SQLITE:
		CreateSqlliteDatabase()
	}

	return nil
}

/*
DeleteMails deletes a set of email and their attachments. If the optional start/end
dates are not empty strings then the delete is date-filtered.
*/
func DeleteMails(startDate, endDate string) error {
	where := ""
	var err error

	if len(startDate) > 0 && len(endDate) > 0 {
		where = where + " AND dateSent >= ? "
		where = where + " AND dateSend <= ? "
	}

	sql := "DELETE FROM mailitem WHERE 1=1" + where

	if len(where) > 0 {
		_, err = golangdb.Db["lib"].Exec(sql)
	} else {
		_, err = golangdb.Db["lib"].Exec(sql, startDate, endDate)
	}

	return err
}

/*
Disconnects from the database storage
*/
func DisconnectFromStorage() {
	golangdb.Db["lib"].Close()
}

/*
Returns an attachment by ID
*/
func GetAttachment(mailId, attachmentId string) (attachment.Attachment, error) {
	result := attachment.Attachment{}

	rows, err := golangdb.Db["lib"].Query(`
		SELECT
			  fileName TEXT
			, contentType TEXT
			, content TEXT
		FROM attachment
		WHERE
			id=?
			AND mailItemId=?
	`, attachmentId, mailId)

	if err != nil {
		return result, fmt.Errorf("Error running query to get attachment")
	}

	defer rows.Close()
	rows.Next()

	var fileName string
	var contentType string
	var content string

	rows.Scan(&fileName, &contentType, &content)

	result.Headers = &attachment.AttachmentHeader{
		FileName:    fileName,
		ContentType: contentType,
	}

	result.Contents = content
	return result, nil
}

func getMailQuery(whereClause string) string {
	sql := `
		SELECT
			  mailitem.id AS mailItemId
			, mailitem.dateSent
			, mailitem.fromAddress
			, mailitem.toAddressList
			, mailitem.subject
			, mailitem.xmailer
			, mailitem.body
			, mailitem.contentType
			, mailitem.boundary

		FROM mailitem

		WHERE 1=1 `

	sql = sql + whereClause
	sql = sql + ` ORDER BY mailitem.dateSent DESC `

	return sql
}

/*
Returns a single mail item by ID.
*/
func GetMail(id string) (mailitem.MailItem, error) {
	result := mailitem.MailItem{}
	xssService := sanitization.NewXSSService()

	sql := getMailQuery(" AND mailitem.id=? ")
	rows, err := golangdb.Db["lib"].Query(sql, id)

	if err != nil {
		return result, fmt.Errorf("Error running query to get mail items: %s", err)
	}

	rows.Next()

	var mailItemId string
	var dateSent string
	var fromAddress string
	var toAddressList string
	var subject string
	var xmailer string
	var body string
	var contentType string
	var boundary string

	rows.Scan(&mailItemId, &dateSent, &fromAddress, &toAddressList, &subject, &xmailer, &body, &contentType, &boundary)

	result = mailitem.MailItem{
		Id:          mailItemId,
		DateSent:    dateSent,
		FromAddress: fromAddress,
		ToAddresses: strings.Split(toAddressList, "; "),
		Subject:     xssService.SanitizeString(subject),
		XMailer:     xssService.SanitizeString(xmailer),
		Body:        xssService.SanitizeString(body),
		ContentType: contentType,
		Boundary:    boundary,
	}

	/*
	 * Get attachments
	 */
	sql = `
		SELECT
			  attachment.id AS attachmentId
			, attachment.fileName
			, attachment.contentType

		FROM attachment
		WHERE
			attachment.mailItemId=?`

	attachmentRows, err := golangdb.Db["lib"].Query(sql, mailItemId)
	if err != nil {
		return result, err
	}

	attachments := make([]*attachment.Attachment, 0)

	for attachmentRows.Next() {
		var attachmentId string
		var fileName string
		var contentType string

		attachmentRows.Scan(&attachmentId, &fileName, &contentType)

		newAttachment := &attachment.Attachment{
			Id:     attachmentId,
			MailId: mailItemId,
			Headers: &attachment.AttachmentHeader{
				FileName:    xssService.SanitizeString(fileName),
				ContentType: contentType,
			},
		}

		attachments = append(attachments, newAttachment)
	}

	attachmentRows.Close()

	result.Attachments = attachments

	rows.Close()
	return result, nil
}

/*
Retrieves all stored mail items as an array of MailItem items. Only
returns rows starting at offset and gets up to length records. NOTE:
This code stinks. It gets ALL rows, then returns a slice and the total number
of mail items. Ick!
*/
func GetMailCollection(offset, length int) ([]mailitem.MailItem, int, error) {
	result := make([]mailitem.MailItem, 0)
	xssService := sanitization.NewXSSService()

	sql := getMailQuery("")
	rows, err := golangdb.Db["lib"].Query(sql)

	if err != nil {
		return result, 0, fmt.Errorf("Error running query to get mail items: %s", err)
	}

	/*
	 * Loop over our records and grab attachments on the way.
	 */
	for rows.Next() {
		var mailItemId string
		var dateSent string
		var fromAddress string
		var toAddressList string
		var subject string
		var xmailer string
		var body string
		var contentType string
		var boundary string

		rows.Scan(&mailItemId, &dateSent, &fromAddress, &toAddressList, &subject, &xmailer, &body, &contentType, &boundary)

		newItem := mailitem.MailItem{
			Id:          mailItemId,
			DateSent:    dateSent,
			FromAddress: fromAddress,
			ToAddresses: strings.Split(toAddressList, "; "),
			Subject:     xssService.SanitizeString(subject),
			XMailer:     xssService.SanitizeString(xmailer),
			Body:        xssService.SanitizeString(body),
			ContentType: contentType,
			Boundary:    boundary,
		}

		/*
		 * Get attachments
		 */
		sql = `
			SELECT
				  attachment.id AS attachmentId
				, attachment.fileName
			FROM attachment
			WHERE
				attachment.mailItemId=?`

		attachmentRows, err := golangdb.Db["lib"].Query(sql, mailItemId)
		if err != nil {
			return result, 0, err
		}

		attachments := make([]*attachment.Attachment, 0)

		for attachmentRows.Next() {
			var attachmentId string
			var fileName string

			attachmentRows.Scan(&attachmentId, &fileName)

			newAttachment := &attachment.Attachment{
				Id:      attachmentId,
				MailId:  mailItemId,
				Headers: &attachment.AttachmentHeader{FileName: xssService.SanitizeString(fileName)},
			}

			attachments = append(attachments, newAttachment)
		}

		attachmentRows.Close()

		newItem.Attachments = attachments
		result = append(result, newItem)
	}

	rows.Close()

	totalRecords := len(result)
	return result, totalRecords, nil
}

/*
GetMailCount returns a count of mail items in the storage system.
*/
func GetMailCount() (int, error) {
	var mailItemCount int
	var err error

	err = golangdb.Db["lib"].QueryRow(`SELECT COUNT(id) AS mailItemCount FROM mailitem`).Scan(&mailItemCount)
	if err != nil {
		return 0, fmt.Errorf("Error running query to get mail item count: %s", err)
	}

	return mailItemCount, nil
}

func storeAttachments(mailItemId string, transaction *sql.Tx, attachments []*attachment.Attachment) error {
	for _, a := range attachments {
		attachmentId, err := mailitem.GenerateId()
		if err != nil {
			return fmt.Errorf("Error generating ID for attachment: %s", err)
		}

		statement, err := transaction.Prepare(`
			INSERT INTO attachment (
				  id
				, mailItemId
				, fileName
				, contentType
				, content
			) VALUES (
				  ?
				, ?
				, ?
				, ?
				, ?
			)
		`)

		if err != nil {
			return fmt.Errorf("Error preparing insert attachment statement: %s", err)
		}

		_, err = statement.Exec(
			attachmentId,
			mailItemId,
			a.Headers.FileName,
			a.Headers.ContentType,
			a.Contents,
		)

		if err != nil {
			return fmt.Errorf("Error executing insert attachment in StoreMail: %s", err)
		}

		statement.Close()
		a.Id = attachmentId
	}

	return nil
}

func StoreMail(mailItem *mailitem.MailItem) (string, error) {
	/*
	 * Create a transaction and insert the new mail item
	 */
	transaction, err := golangdb.Db["lib"].Begin()
	if err != nil {
		return "", fmt.Errorf("Error starting transaction in StoreMail: %s", err)
	}

	/*
	 * Insert the mail item
	 */
	statement, err := transaction.Prepare(`
			INSERT INTO mailitem (
				  id
				, dateSent
				, fromAddress
				, toAddressList
				, subject
				, xmailer
				, body
				, contentType
				, boundary
			) VALUES (
				  ?
				, ?
				, ?
				, ?
				, ?
				, ?
				, ?
				, ?
				, ?
			)
		`)

	if err != nil {
		return "", fmt.Errorf("Error preparing insert statement for mail item in StoreMail: %s", err)
	}

	_, err = statement.Exec(
		mailItem.Id,
		mailItem.DateSent,
		mailItem.FromAddress,
		strings.Join(mailItem.ToAddresses, "; "),
		mailItem.Subject,
		mailItem.XMailer,
		mailItem.Body,
		mailItem.ContentType,
		mailItem.Boundary,
	)

	if err != nil {
		transaction.Rollback()
		return "", fmt.Errorf("Error executing insert for mail item in StoreMail: %s", err)
	}

	statement.Close()

	/*
	 * Insert attachments
	 */
	if err = storeAttachments(mailItem.Id, transaction, mailItem.Attachments); err != nil {
		transaction.Rollback()
		return "", fmt.Errorf("Unable to insert attachments in StoreMail: %s", err)
	}

	transaction.Commit()
	log.Printf("New mail item written to database.\n\n")

	return mailItem.Id, nil
}
<|endoftext|>


<commit_before>package storage

import (
	"fmt"
	"github.com/viant/toolbox"
	"io"
	"net/url"
	"path"
)

//Service represents abstract way to accessing local or remote storage
type Service interface {
	//List returns a list of object for supplied url
	List(URL string) ([]Object, error)

	//Exists returns true if resource exists
	Exists(URL string) (bool, error)

	//Object returns a Object for supplied url
	StorageObject(URL string) (Object, error)

	//Download returns reader for downloaded storage object
	Download(object Object) (io.Reader, error)

	//Upload uploads provided reader content for supplied storage object.
	Upload(URL string, reader io.Reader) error

	//Delete removes passed in storage object
	Delete(object Object) error

	//Register register schema with provided service
	Register(schema string, service Service) error
}

type storageService struct {
	registry map[string]Service
}

func (s *storageService) getServiceForSchema(URL string) (Service, error) {
	parsedUrl, err := url.Parse(URL)
	if err != nil {
		return nil, err
	}

	if result, found := s.registry[parsedUrl.Scheme]; found {
		return result, nil
	}

	return nil, fmt.Errorf("Failed to lookup url schema %v in %v", parsedUrl.Scheme, URL)
}

//List lists all object for passed in URL
func (s *storageService) List(URL string) ([]Object, error) {
	service, err := s.getServiceForSchema(URL)
	if err != nil {
		return nil, err
	}
	return service.List(URL)
}

//Exists returns true if resource exists
func (s *storageService) Exists(URL string) (bool, error) {
	service, err := s.getServiceForSchema(URL)
	if err != nil {
		return false, err
	}
	return service.Exists(URL)
}

//StorageObject returns storage object for provided URL
func (s *storageService) StorageObject(URL string) (Object, error) {
	service, err := s.getServiceForSchema(URL)
	if err != nil {
		return nil, err
	}
	return service.StorageObject(URL)
}

//Download downloads content for passed in object
func (s *storageService) Download(object Object) (io.Reader, error) {
	service, err := s.getServiceForSchema(object.URL())
	if err != nil {
		return nil, err
	}
	return service.Download(object)
}

//Uploads content for passed in URL
func (s *storageService) Upload(URL string, reader io.Reader) error {
	service, err := s.getServiceForSchema(URL)
	if err != nil {
		return err
	}
	return service.Upload(URL, reader)
}

//Delete remove storage object
func (s *storageService) Delete(object Object) error {
	service, err := s.getServiceForSchema(object.URL())
	if err != nil {
		return err
	}
	return service.Delete(object)
}

//Register register storage schema
func (s *storageService) Register(schema string, service Service) error {
	s.registry[schema] = service
	return nil
}

//NewService creates a new storage service
func NewService() Service {
	var result = &storageService{
		registry: make(map[string]Service),
	}
	result.Register("file", &fileStorageService{})
	return result
}

//NewServiceForURL creates a new storage service for provided URL scheme and optional credential file
func NewServiceForURL(URL, credentialFile string) (Service, error) {
	parsedURL, err := url.Parse(URL)
	if err != nil {
		return nil, err
	}
	service := NewService()
	provider := NewStorageProvider().Get(parsedURL.Scheme)

	if provider != nil {
		serviceForScheme, err := provider(credentialFile)
		if err != nil {
			return nil, fmt.Errorf("Failed to get storage for url %v: %v", URL, err)
		}
		err = service.Register(parsedURL.Scheme, serviceForScheme)
		if err != nil {
			return nil, err
		}
	} else if parsedURL.Scheme != "file" {
		return nil, fmt.Errorf("Unsupported scheme %v", URL)
	}
	return service, nil
}

func copy(sourceService Service, sourceURL string, targetService Service, targetURL string, modifyContentHandler func(reader io.Reader) (io.Reader, error), subPath string) error {
	sourceListURL := sourceURL
	if subPath != "" {
		sourceListURL = toolbox.URLPathJoin(sourceURL, subPath)
	}
	objects, err := sourceService.List(sourceListURL)
	var objectRelativePath string
	for _, object := range objects {

		if object.URL() == sourceURL && object.IsFolder() {
			continue
		}

		if len(object.URL()) > len(sourceURL) {
			objectRelativePath = object.URL()[len(sourceURL):]
		}
		var targetObjectURL = targetURL
		if objectRelativePath != "" {
			targetObjectURL = toolbox.URLPathJoin(targetURL, objectRelativePath)
		}
		var reader io.Reader
		if object.IsContent() {
			reader, err = sourceService.Download(object)
			if err != nil {
				err = fmt.Errorf("Unable download, %v -> %v, %v", object.URL(), targetObjectURL, err)
				return err
			}

			if modifyContentHandler != nil {
				reader, err = modifyContentHandler(reader)
				if err != nil {
					err = fmt.Errorf("Unable modify content, %v %v %v", object.URL(), targetObjectURL, err)
					return err
				}
			}

			targetObject, err := targetService.StorageObject(targetObjectURL)
			if (targetObject != nil && targetObject.IsFolder()) {
				_, file := path.Split(object.URL())
				targetObjectURL = toolbox.URLPathJoin(targetObjectURL, file)
			}

			err = targetService.Upload(targetObjectURL, reader)
			if err != nil {
				err = fmt.Errorf("Unable upload, %v %v %v", object.URL(), targetObjectURL, err)
				return err
			}
		} else {
			err = copy(sourceService, sourceURL, targetService, targetURL, modifyContentHandler, objectRelativePath)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

//Copy downloads objects from source URL to upload them to target URL.
func Copy(sourceService Service, sourceURL string, targetService Service, targetURL string, modifyContentHandler func(reader io.Reader) (io.Reader, error)) (err error) {
	err = copy(sourceService, sourceURL, targetService, targetURL, modifyContentHandler, "")
	if err != nil {
		err = fmt.Errorf("Failed to copy %v -> %v: %v", sourceURL, targetURL, err)
	}
	return err
}
<commit_msg>updated copy function to enable custom copy handler<commit_after>package storage

import (
	"fmt"
	"github.com/viant/toolbox"
	"io"
	"net/url"
	"path"
)

//Service represents abstract way to accessing local or remote storage
type Service interface {
	//List returns a list of object for supplied url
	List(URL string) ([]Object, error)

	//Exists returns true if resource exists
	Exists(URL string) (bool, error)

	//Object returns a Object for supplied url
	StorageObject(URL string) (Object, error)

	//Download returns reader for downloaded storage object
	Download(object Object) (io.Reader, error)

	//Upload uploads provided reader content for supplied storage object.
	Upload(URL string, reader io.Reader) error

	//Delete removes passed in storage object
	Delete(object Object) error

	//Register register schema with provided service
	Register(schema string, service Service) error
}

type CopyHandler func(sourceObject Object, source io.Reader, destinationService Service, destinationURL string) error
type ModificationHandler func(reader io.Reader) (io.Reader, error)

type storageService struct {
	registry map[string]Service
}

func (s *storageService) getServiceForSchema(URL string) (Service, error) {
	parsedUrl, err := url.Parse(URL)
	if err != nil {
		return nil, err
	}

	if result, found := s.registry[parsedUrl.Scheme]; found {
		return result, nil
	}

	return nil, fmt.Errorf("Failed to lookup url schema %v in %v", parsedUrl.Scheme, URL)
}

//List lists all object for passed in URL
func (s *storageService) List(URL string) ([]Object, error) {
	service, err := s.getServiceForSchema(URL)
	if err != nil {
		return nil, err
	}
	return service.List(URL)
}

//Exists returns true if resource exists
func (s *storageService) Exists(URL string) (bool, error) {
	service, err := s.getServiceForSchema(URL)
	if err != nil {
		return false, err
	}
	return service.Exists(URL)
}

//StorageObject returns storage object for provided URL
func (s *storageService) StorageObject(URL string) (Object, error) {
	service, err := s.getServiceForSchema(URL)
	if err != nil {
		return nil, err
	}
	return service.StorageObject(URL)
}

//Download downloads content for passed in object
func (s *storageService) Download(object Object) (io.Reader, error) {
	service, err := s.getServiceForSchema(object.URL())
	if err != nil {
		return nil, err
	}
	return service.Download(object)
}

//Uploads content for passed in URL
func (s *storageService) Upload(URL string, reader io.Reader) error {
	service, err := s.getServiceForSchema(URL)
	if err != nil {
		return err
	}
	return service.Upload(URL, reader)
}

//Delete remove storage object
func (s *storageService) Delete(object Object) error {
	service, err := s.getServiceForSchema(object.URL())
	if err != nil {
		return err
	}
	return service.Delete(object)
}

//Register register storage schema
func (s *storageService) Register(schema string, service Service) error {
	s.registry[schema] = service
	return nil
}

//NewService creates a new storage service
func NewService() Service {
	var result = &storageService{
		registry: make(map[string]Service),
	}
	result.Register("file", &fileStorageService{})
	return result
}

//NewServiceForURL creates a new storage service for provided URL scheme and optional credential file
func NewServiceForURL(URL, credentialFile string) (Service, error) {
	parsedURL, err := url.Parse(URL)
	if err != nil {
		return nil, err
	}
	service := NewService()
	provider := NewStorageProvider().Get(parsedURL.Scheme)

	if provider != nil {
		serviceForScheme, err := provider(credentialFile)
		if err != nil {
			return nil, fmt.Errorf("Failed to get storage for url %v: %v", URL, err)
		}
		err = service.Register(parsedURL.Scheme, serviceForScheme)
		if err != nil {
			return nil, err
		}
	} else if parsedURL.Scheme != "file" {
		return nil, fmt.Errorf("Unsupported scheme %v", URL)
	}
	return service, nil
}

func copy(sourceService Service, sourceURL string, destinationService Service, destinationURL string, modifyContentHandler ModificationHandler, subPath string, copyHandler CopyHandler) error {
	sourceListURL := sourceURL
	if subPath != "" {
		sourceListURL = toolbox.URLPathJoin(sourceURL, subPath)
	}
	objects, err := sourceService.List(sourceListURL)
	var objectRelativePath string
	for _, object := range objects {
		if object.URL() == sourceURL && object.IsFolder() {
			continue
		}
		if len(object.URL()) > len(sourceURL) {
			objectRelativePath = object.URL()[len(sourceURL):]
		}
		var destinationObjectURL = destinationURL
		if objectRelativePath != "" {
			destinationObjectURL = toolbox.URLPathJoin(destinationURL, objectRelativePath)
		}
		var reader io.Reader
		if object.IsContent() {
			reader, err = sourceService.Download(object)
			if err != nil {
				err = fmt.Errorf("Unable download, %v -> %v, %v", object.URL(), destinationObjectURL, err)
				return err
			}

			if modifyContentHandler != nil {
				reader, err = modifyContentHandler(reader)
				if err != nil {
					err = fmt.Errorf("Unable modify content, %v %v %v", object.URL(), destinationObjectURL, err)
					return err
				}
			}
			destinationObject, err := destinationService.StorageObject(destinationObjectURL)
			if (subPath == "" && destinationObject != nil && destinationObject.IsFolder()) {
				_, file := path.Split(object.URL())
				destinationObjectURL = toolbox.URLPathJoin(destinationObjectURL, file)
			}
			err = copyHandler(object, reader, destinationService, destinationObjectURL)
			if err != nil {
				return nil
			}

		} else {
			err = copy(sourceService, sourceURL, destinationService, destinationURL, modifyContentHandler, objectRelativePath, copyHandler)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func copySourceToDestination(sourceObject Object, reader io.Reader, destinationService Service, destinationURL string) error {
	err := destinationService.Upload(destinationURL, reader)
	if err != nil {
		err = fmt.Errorf("Unable upload, %v %v %v", sourceObject.URL(), destinationURL, err)

	}
	return err
}

//Copy downloads objects from source URL to upload them to destination URL.
func Copy(sourceService Service, sourceURL string, destinationService Service, destinationURL string, modifyContentHandler ModificationHandler, copyHandler CopyHandler) (err error) {
	if copyHandler == nil {
		copyHandler = copySourceToDestination
	}
	err = copy(sourceService, sourceURL, destinationService, destinationURL, modifyContentHandler, "", copyHandler)
	if err != nil {
		err = fmt.Errorf("Failed to copy %v -> %v: %v", sourceURL, destinationURL, err)
	}
	return err
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 storage is a Google Cloud Storage client.
package storage

import (
	"errors"
	"fmt"
	"io"
	"net/http"

	raw "code.google.com/p/google-api-go-client/storage/v1"
)

const (
	// ScopeFullControl grants permissions to manage your
	// data and permissions in Google Cloud Storage.
	ScopeFullControl = raw.DevstorageFull_controlScope

	// ScopeReadOnly grants permissions to
	// view your data in Google Cloud Storage.
	ScopeReadOnly = raw.DevstorageRead_onlyScope

	// ScopeReadWrite grants permissions to manage your
	// data in Google Cloud Storage.
	ScopeReadWrite = raw.DevstorageRead_writeScope
)

const (
	templUrlMedia = "https://storage.googleapis.com/%s/%s"
)

// Bucket represents a Google Cloud Storage bucket.
type Bucket struct {
	// Name is the name of the bucket.
	Name string `json:"name,omitempty"`
}

type conn struct {
	c *http.Client
	s *raw.Service
}

// BucketClient is a client to perform object operations on.
type BucketClient struct {
	name string
	conn *conn
}

// String returns a string representation of the bucket client.
// E.g. <bucket: my-project-bucket>
func (b *BucketClient) String() string {
	return fmt.Sprintf("<bucket: %v>", b.name)
}

// Client represents a Google Cloud Storage client.
type Client struct {
	conn *conn
}

// New returns a new Google Cloud Storage client. The provided
// RoundTripper should be authorized and authenticated to make
// calls to Google Cloud Storage API.
func New(tr http.RoundTripper) *Client {
	return NewWithClient(&http.Client{Transport: tr})
}

// NewWithClient returns a new Google Cloud Storage client that
// uses the provided http.Client. Provided http.Client is responsible
// to authorize and authenticate the requests made to the
// Google Cloud Storage API.
func NewWithClient(c *http.Client) *Client {
	s, _ := raw.New(c)
	return &Client{conn: &conn{s: s, c: c}}
}

// TODO(jbd): Add storage.buckets.list.
// TODO(jbd): Add storage.buckets.insert.
// TODO(jbd): Add storage.buckets.update.
// TODO(jbd): Add storage.buckets.delete.

// TODO(jbd): Add storage.objects.watch.

// Bucket returns the metadata for the specified bucket.
func (c *Client) Bucket(name string) (*Bucket, error) {
	panic("not yet implemented")
}

// BucketClient returns a bucket client to perform object operations on.
func (c *Client) BucketClient(bucketname string) *BucketClient {
	return &BucketClient{name: bucketname, conn: c.conn}
}

// List lists objects from the bucket. You can specify a query
// to filter the results. If q is nil, no filtering is applied.
func (b *BucketClient) List(q *Query) (*Objects, error) {
	c := b.conn.s.Objects.List(b.name)
	if q != nil {
		c.Delimiter(q.Delimiter)
		c.Prefix(q.Prefix)
		c.Versions(q.Versions)
		c.PageToken(q.Cursor)
		if q.MaxResults > 0 {
			c.MaxResults(int64(q.MaxResults))
		}
	}
	resp, err := c.Do()
	if err != nil {
		return nil, err
	}
	objects := &Objects{
		Results: make([]*Object, len(resp.Items)),
	}
	for i, item := range resp.Items {
		objects.Results[i] = newObject(item)
	}
	if resp.NextPageToken != "" {
		next := Query{}
		if q != nil {
			// keep the other filtering
			// criteria if there is a query
			next = *q
		}
		next.Cursor = resp.NextPageToken
		objects.Next = &next
	}
	return objects, nil
}

// Stat returns meta information about the specified object.
func (b *BucketClient) Stat(name string) (*Object, error) {
	o, err := b.conn.s.Objects.Get(b.name, name).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// Put inserts/updates an object with the provided meta information.
func (b *BucketClient) Put(name string, info *Object) (*Object, error) {
	o, err := b.conn.s.Objects.Insert(b.name, info.toRawObject()).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// Delete deletes the specified object.
func (b *BucketClient) Delete(name string) error {
	return b.conn.s.Objects.Delete(b.name, name).Do()
}

// Copy copies the source object to the destination with the new
// meta information provided.
// The destination object is inserted into the source bucket
// if the destination object doesn't specify another bucket name.
func (b *BucketClient) Copy(name string, dest *Object) (*Object, error) {
	if dest.Name == "" {
		return nil, errors.New("storage: missing dest name")
	}
	destBucket := dest.Bucket
	if destBucket == "" {
		destBucket = b.name
	}
	o, err := b.conn.s.Objects.Copy(
		b.name, name, destBucket, dest.Name, dest.toRawObject()).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// NewReader creates a new io.ReadCloser to read the contents
// of the object.
func (b *BucketClient) NewReader(name string) (io.ReadCloser, error) {
	resp, err := b.conn.c.Get(fmt.Sprintf(templUrlMedia, b.name, name))
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

// NewWriter returns a new ObjectWriter to write to the GCS object
// identified by the specified object name.
// If such object doesn't exist, it creates one. If info is not nil,
// write operation also modifies the meta information of the object.
// All read-only fields are ignored during metadata updates.
func (b *BucketClient) NewWriter(name string, info *Object) *ObjectWriter {
	i := Object{}
	if info != nil {
		i = *info
	}
	i.Bucket = b.name
	i.Name = name
	return newObjectWriter(b.conn, &i)
}
<commit_msg>Add support for directory-mode listings by returning the Prefixes from a List call.<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 storage is a Google Cloud Storage client.
package storage

import (
	"errors"
	"fmt"
	"io"
	"net/http"

	raw "code.google.com/p/google-api-go-client/storage/v1"
)

const (
	// ScopeFullControl grants permissions to manage your
	// data and permissions in Google Cloud Storage.
	ScopeFullControl = raw.DevstorageFull_controlScope

	// ScopeReadOnly grants permissions to
	// view your data in Google Cloud Storage.
	ScopeReadOnly = raw.DevstorageRead_onlyScope

	// ScopeReadWrite grants permissions to manage your
	// data in Google Cloud Storage.
	ScopeReadWrite = raw.DevstorageRead_writeScope
)

const (
	templUrlMedia = "https://storage.googleapis.com/%s/%s"
)

// Bucket represents a Google Cloud Storage bucket.
type Bucket struct {
	// Name is the name of the bucket.
	Name string `json:"name,omitempty"`
}

type conn struct {
	c *http.Client
	s *raw.Service
}

// BucketClient is a client to perform object operations on.
type BucketClient struct {
	name string
	conn *conn
}

// String returns a string representation of the bucket client.
// E.g. <bucket: my-project-bucket>
func (b *BucketClient) String() string {
	return fmt.Sprintf("<bucket: %v>", b.name)
}

// Client represents a Google Cloud Storage client.
type Client struct {
	conn *conn
}

// New returns a new Google Cloud Storage client. The provided
// RoundTripper should be authorized and authenticated to make
// calls to Google Cloud Storage API.
func New(tr http.RoundTripper) *Client {
	return NewWithClient(&http.Client{Transport: tr})
}

// NewWithClient returns a new Google Cloud Storage client that
// uses the provided http.Client. Provided http.Client is responsible
// to authorize and authenticate the requests made to the
// Google Cloud Storage API.
func NewWithClient(c *http.Client) *Client {
	s, _ := raw.New(c)
	return &Client{conn: &conn{s: s, c: c}}
}

// TODO(jbd): Add storage.buckets.list.
// TODO(jbd): Add storage.buckets.insert.
// TODO(jbd): Add storage.buckets.update.
// TODO(jbd): Add storage.buckets.delete.

// TODO(jbd): Add storage.objects.watch.

// Bucket returns the metadata for the specified bucket.
func (c *Client) Bucket(name string) (*Bucket, error) {
	panic("not yet implemented")
}

// BucketClient returns a bucket client to perform object operations on.
func (c *Client) BucketClient(bucketname string) *BucketClient {
	return &BucketClient{name: bucketname, conn: c.conn}
}

// List lists objects from the bucket. You can specify a query
// to filter the results. If q is nil, no filtering is applied.
func (b *BucketClient) List(q *Query) (*Objects, error) {
	c := b.conn.s.Objects.List(b.name)
	if q != nil {
		c.Delimiter(q.Delimiter)
		c.Prefix(q.Prefix)
		c.Versions(q.Versions)
		c.PageToken(q.Cursor)
		if q.MaxResults > 0 {
			c.MaxResults(int64(q.MaxResults))
		}
	}
	resp, err := c.Do()
	if err != nil {
		return nil, err
	}
	objects := &Objects{
		Results:  make([]*Object, len(resp.Items)),
		Prefixes: make([]string, len(resp.Prefixes)),
	}
	for i, item := range resp.Items {
		objects.Results[i] = newObject(item)
	}
	for i, prefix := range resp.Prefixes {
		objects.Prefixes[i] = prefix
	}
	if resp.NextPageToken != "" {
		next := Query{}
		if q != nil {
			// keep the other filtering
			// criteria if there is a query
			next = *q
		}
		next.Cursor = resp.NextPageToken
		objects.Next = &next
	}
	return objects, nil
}

// Stat returns meta information about the specified object.
func (b *BucketClient) Stat(name string) (*Object, error) {
	o, err := b.conn.s.Objects.Get(b.name, name).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// Put inserts/updates an object with the provided meta information.
func (b *BucketClient) Put(name string, info *Object) (*Object, error) {
	o, err := b.conn.s.Objects.Insert(b.name, info.toRawObject()).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// Delete deletes the specified object.
func (b *BucketClient) Delete(name string) error {
	return b.conn.s.Objects.Delete(b.name, name).Do()
}

// Copy copies the source object to the destination with the new
// meta information provided.
// The destination object is inserted into the source bucket
// if the destination object doesn't specify another bucket name.
func (b *BucketClient) Copy(name string, dest *Object) (*Object, error) {
	if dest.Name == "" {
		return nil, errors.New("storage: missing dest name")
	}
	destBucket := dest.Bucket
	if destBucket == "" {
		destBucket = b.name
	}
	o, err := b.conn.s.Objects.Copy(
		b.name, name, destBucket, dest.Name, dest.toRawObject()).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// NewReader creates a new io.ReadCloser to read the contents
// of the object.
func (b *BucketClient) NewReader(name string) (io.ReadCloser, error) {
	resp, err := b.conn.c.Get(fmt.Sprintf(templUrlMedia, b.name, name))
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

// NewWriter returns a new ObjectWriter to write to the GCS object
// identified by the specified object name.
// If such object doesn't exist, it creates one. If info is not nil,
// write operation also modifies the meta information of the object.
// All read-only fields are ignored during metadata updates.
func (b *BucketClient) NewWriter(name string, info *Object) *ObjectWriter {
	i := Object{}
	if info != nil {
		i = *info
	}
	i.Bucket = b.name
	i.Name = name
	return newObjectWriter(b.conn, &i)
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 storage is a Google Cloud Storage client.
package storage

import (
	"errors"
	"fmt"
	"io"
	"net/http"

	"code.google.com/p/google-api-go-client/googleapi"
	raw "code.google.com/p/google-api-go-client/storage/v1"
	"google.golang.org/cloud/internal"
)

var (
	ErrBucketNotExists = errors.New("storage: bucket doesn't exist")
	ErrObjectNotExists = errors.New("storage: object doesn't exist")
)

const (
	// ScopeFullControl grants permissions to manage your
	// data and permissions in Google Cloud Storage.
	ScopeFullControl = raw.DevstorageFull_controlScope

	// ScopeReadOnly grants permissions to
	// view your data in Google Cloud Storage.
	ScopeReadOnly = raw.DevstorageRead_onlyScope

	// ScopeReadWrite grants permissions to manage your
	// data in Google Cloud Storage.
	ScopeReadWrite = raw.DevstorageRead_writeScope
)

const (
	templURLMedia = "https://storage.googleapis.com/%s/%s"
)

type conn struct {
	c *http.Client
	s *raw.Service
}

// BucketClient is a client to perform object operations on.
type BucketClient struct {
	name string
	conn *conn
}

// String returns a string representation of the bucket client.
// E.g. <bucket: my-project-bucket>
func (b *BucketClient) String() string {
	return fmt.Sprintf("<bucket: %v>", b.name)
}

// Client represents a Google Cloud Storage client.
type Client struct {
	projID string
	conn   *conn
}

// New returns a new Google Cloud Storage client. The provided
// RoundTripper should be authorized and authenticated to make
// calls to Google Cloud Storage API.
// You can obtain the project ID from the Google Developers Console,
// https://console.developers.google.com.
func New(projID string, tr http.RoundTripper) *Client {
	return NewWithClient(projID, &http.Client{Transport: tr})
}

// NewWithClient returns a new Google Cloud Storage client that
// uses the provided http.Client. Provided http.Client is responsible
// to authorize and authenticate the requests made to the
// Google Cloud Storage API.
// It mutates the client's original Transport to append the cloud
// package's user-agent to the outgoing requests.
// You can obtain the project ID from the Google Developers Console,
// https://console.developers.google.com.
func NewWithClient(projID string, c *http.Client) *Client {
	c.Transport = &internal.UATransport{Base: c.Transport}
	s, _ := raw.New(c)
	return &Client{projID: projID, conn: &conn{s: s, c: c}}
}

// TODO(jbd): Add storage.buckets.list.
// TODO(jbd): Add storage.buckets.insert.
// TODO(jbd): Add storage.buckets.update.
// TODO(jbd): Add storage.buckets.delete.

// TODO(jbd): Add storage.objects.watch.

// Bucket returns the metadata for the specified bucket.
func (c *Client) Bucket(name string) (*Bucket, error) {
	resp, err := c.conn.s.Buckets.Get(name).Do()
	if e, ok := err.(*googleapi.Error); ok && e.Code == http.StatusNotFound {
		return nil, ErrBucketNotExists
	}
	if err != nil {
		return nil, err
	}
	return newBucket(resp), nil
}

// BucketClient returns a bucket client to perform object operations on.
func (c *Client) BucketClient(bucketname string) *BucketClient {
	return &BucketClient{name: bucketname, conn: c.conn}
}

// List lists objects from the bucket. You can specify a query
// to filter the results. If q is nil, no filtering is applied.
func (b *BucketClient) List(q *Query) (*Objects, error) {
	c := b.conn.s.Objects.List(b.name)
	if q != nil {
		c.Delimiter(q.Delimiter)
		c.Prefix(q.Prefix)
		c.Versions(q.Versions)
		c.PageToken(q.Cursor)
		if q.MaxResults > 0 {
			c.MaxResults(int64(q.MaxResults))
		}
	}
	resp, err := c.Do()
	if err != nil {
		return nil, err
	}
	objects := &Objects{
		Results:  make([]*Object, len(resp.Items)),
		Prefixes: make([]string, len(resp.Prefixes)),
	}
	for i, item := range resp.Items {
		objects.Results[i] = newObject(item)
	}
	for i, prefix := range resp.Prefixes {
		objects.Prefixes[i] = prefix
	}
	if resp.NextPageToken != "" {
		next := Query{}
		if q != nil {
			// keep the other filtering
			// criteria if there is a query
			next = *q
		}
		next.Cursor = resp.NextPageToken
		objects.Next = &next
	}
	return objects, nil
}

// Stat returns meta information about the specified object.
func (b *BucketClient) Stat(name string) (*Object, error) {
	o, err := b.conn.s.Objects.Get(b.name, name).Do()
	if e, ok := err.(*googleapi.Error); ok && e.Code == http.StatusNotFound {
		return nil, ErrObjectNotExists
	}
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// Put inserts/updates an object with the provided meta information.
func (b *BucketClient) Put(name string, info *Object) (*Object, error) {
	o, err := b.conn.s.Objects.Insert(b.name, info.toRawObject()).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// Delete deletes the specified object.
func (b *BucketClient) Delete(name string) error {
	return b.conn.s.Objects.Delete(b.name, name).Do()
}

// Copy copies the source object to the destination with the new
// meta information provided.
// The destination object is inserted into the source bucket
// if the destination object doesn't specify another bucket name.
func (b *BucketClient) Copy(name string, dest *Object) (*Object, error) {
	if dest.Name == "" {
		return nil, errors.New("storage: missing dest name")
	}
	if dest.Bucket == "" {
		// Make a copy of the dest object instead of mutating it.
		dest2 := *dest
		dest2.Bucket = b.name
		dest = &dest2
	}
	o, err := b.conn.s.Objects.Copy(
		b.name, name, dest.Bucket, dest.Name, dest.toRawObject()).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// NewReader creates a new io.ReadCloser to read the contents
// of the object.
func (b *BucketClient) NewReader(name string) (io.ReadCloser, error) {
	resp, err := b.conn.c.Get(fmt.Sprintf(templUrlMedia, b.name, name))
	if e, ok := err.(*googleapi.Error); ok && e.Code == http.StatusNotFound {
		return nil, ErrObjectNotExists
	}
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

// NewWriter returns a new ObjectWriter to write to the GCS object
// identified by the specified object name.
// If such object doesn't exist, it creates one. If info is not nil,
// write operation also modifies the meta information of the object.
// All read-only fields are ignored during metadata updates.
func (b *BucketClient) NewWriter(name string, info *Object) *ObjectWriter {
	i := Object{}
	if info != nil {
		i = *info
	}
	i.Bucket = b.name
	i.Name = name
	return newObjectWriter(b.conn, &i)
}
<commit_msg>Fixing the broken build.<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 storage is a Google Cloud Storage client.
package storage

import (
	"errors"
	"fmt"
	"io"
	"net/http"

	"code.google.com/p/google-api-go-client/googleapi"
	raw "code.google.com/p/google-api-go-client/storage/v1"
	"google.golang.org/cloud/internal"
)

var (
	ErrBucketNotExists = errors.New("storage: bucket doesn't exist")
	ErrObjectNotExists = errors.New("storage: object doesn't exist")
)

const (
	// ScopeFullControl grants permissions to manage your
	// data and permissions in Google Cloud Storage.
	ScopeFullControl = raw.DevstorageFull_controlScope

	// ScopeReadOnly grants permissions to
	// view your data in Google Cloud Storage.
	ScopeReadOnly = raw.DevstorageRead_onlyScope

	// ScopeReadWrite grants permissions to manage your
	// data in Google Cloud Storage.
	ScopeReadWrite = raw.DevstorageRead_writeScope
)

const (
	templURLMedia = "https://storage.googleapis.com/%s/%s"
)

type conn struct {
	c *http.Client
	s *raw.Service
}

// BucketClient is a client to perform object operations on.
type BucketClient struct {
	name string
	conn *conn
}

// String returns a string representation of the bucket client.
// E.g. <bucket: my-project-bucket>
func (b *BucketClient) String() string {
	return fmt.Sprintf("<bucket: %v>", b.name)
}

// Client represents a Google Cloud Storage client.
type Client struct {
	projID string
	conn   *conn
}

// New returns a new Google Cloud Storage client. The provided
// RoundTripper should be authorized and authenticated to make
// calls to Google Cloud Storage API.
// You can obtain the project ID from the Google Developers Console,
// https://console.developers.google.com.
func New(projID string, tr http.RoundTripper) *Client {
	return NewWithClient(projID, &http.Client{Transport: tr})
}

// NewWithClient returns a new Google Cloud Storage client that
// uses the provided http.Client. Provided http.Client is responsible
// to authorize and authenticate the requests made to the
// Google Cloud Storage API.
// It mutates the client's original Transport to append the cloud
// package's user-agent to the outgoing requests.
// You can obtain the project ID from the Google Developers Console,
// https://console.developers.google.com.
func NewWithClient(projID string, c *http.Client) *Client {
	c.Transport = &internal.UATransport{Base: c.Transport}
	s, _ := raw.New(c)
	return &Client{projID: projID, conn: &conn{s: s, c: c}}
}

// TODO(jbd): Add storage.buckets.list.
// TODO(jbd): Add storage.buckets.insert.
// TODO(jbd): Add storage.buckets.update.
// TODO(jbd): Add storage.buckets.delete.

// TODO(jbd): Add storage.objects.watch.

// Bucket returns the metadata for the specified bucket.
func (c *Client) Bucket(name string) (*Bucket, error) {
	resp, err := c.conn.s.Buckets.Get(name).Do()
	if e, ok := err.(*googleapi.Error); ok && e.Code == http.StatusNotFound {
		return nil, ErrBucketNotExists
	}
	if err != nil {
		return nil, err
	}
	return newBucket(resp), nil
}

// BucketClient returns a bucket client to perform object operations on.
func (c *Client) BucketClient(bucketname string) *BucketClient {
	return &BucketClient{name: bucketname, conn: c.conn}
}

// List lists objects from the bucket. You can specify a query
// to filter the results. If q is nil, no filtering is applied.
func (b *BucketClient) List(q *Query) (*Objects, error) {
	c := b.conn.s.Objects.List(b.name)
	if q != nil {
		c.Delimiter(q.Delimiter)
		c.Prefix(q.Prefix)
		c.Versions(q.Versions)
		c.PageToken(q.Cursor)
		if q.MaxResults > 0 {
			c.MaxResults(int64(q.MaxResults))
		}
	}
	resp, err := c.Do()
	if err != nil {
		return nil, err
	}
	objects := &Objects{
		Results:  make([]*Object, len(resp.Items)),
		Prefixes: make([]string, len(resp.Prefixes)),
	}
	for i, item := range resp.Items {
		objects.Results[i] = newObject(item)
	}
	for i, prefix := range resp.Prefixes {
		objects.Prefixes[i] = prefix
	}
	if resp.NextPageToken != "" {
		next := Query{}
		if q != nil {
			// keep the other filtering
			// criteria if there is a query
			next = *q
		}
		next.Cursor = resp.NextPageToken
		objects.Next = &next
	}
	return objects, nil
}

// Stat returns meta information about the specified object.
func (b *BucketClient) Stat(name string) (*Object, error) {
	o, err := b.conn.s.Objects.Get(b.name, name).Do()
	if e, ok := err.(*googleapi.Error); ok && e.Code == http.StatusNotFound {
		return nil, ErrObjectNotExists
	}
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// Put inserts/updates an object with the provided meta information.
func (b *BucketClient) Put(name string, info *Object) (*Object, error) {
	o, err := b.conn.s.Objects.Insert(b.name, info.toRawObject()).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// Delete deletes the specified object.
func (b *BucketClient) Delete(name string) error {
	return b.conn.s.Objects.Delete(b.name, name).Do()
}

// Copy copies the source object to the destination with the new
// meta information provided.
// The destination object is inserted into the source bucket
// if the destination object doesn't specify another bucket name.
func (b *BucketClient) Copy(name string, dest *Object) (*Object, error) {
	if dest.Name == "" {
		return nil, errors.New("storage: missing dest name")
	}
	if dest.Bucket == "" {
		// Make a copy of the dest object instead of mutating it.
		dest2 := *dest
		dest2.Bucket = b.name
		dest = &dest2
	}
	o, err := b.conn.s.Objects.Copy(
		b.name, name, dest.Bucket, dest.Name, dest.toRawObject()).Do()
	if err != nil {
		return nil, err
	}
	return newObject(o), nil
}

// NewReader creates a new io.ReadCloser to read the contents
// of the object.
func (b *BucketClient) NewReader(name string) (io.ReadCloser, error) {
	resp, err := b.conn.c.Get(fmt.Sprintf(templURLMedia, b.name, name))
	if e, ok := err.(*googleapi.Error); ok && e.Code == http.StatusNotFound {
		return nil, ErrObjectNotExists
	}
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

// NewWriter returns a new ObjectWriter to write to the GCS object
// identified by the specified object name.
// If such object doesn't exist, it creates one. If info is not nil,
// write operation also modifies the meta information of the object.
// All read-only fields are ignored during metadata updates.
func (b *BucketClient) NewWriter(name string, info *Object) *ObjectWriter {
	i := Object{}
	if info != nil {
		i = *info
	}
	i.Bucket = b.name
	i.Name = name
	return newObjectWriter(b.conn, &i)
}
<|endoftext|>


<commit_before>/*
	Package storage provides a unified interface to a number of storage engines.
	Since each storage engine has different capabilities, this package defines a
	number of interfaces and the Engine interface provides a way to query
	which interfaces are implemented by a given storage engine.

	Initially we are concentrating on key-value storage engines but expect to
	expand support to graph and relational databases.

	Each storage engine must implement the following:

		NewStore(path string, create bool, options *Options) (Engine, error)
*/
package storage

import (
	"sync"

	"github.com/janelia-flyem/dvid/dvid"
)

func Shutdown() {
	ShutdownFUSE()
}

// ChunkOp is a type-specific operation with an optional WaitGroup to
// sync mapping before reduce.
type ChunkOp struct {
	Op interface{}
	Wg *sync.WaitGroup
}

// Chunk is the unit passed down channels to chunk handlers.  Chunks can be passed
// from lower-level database access functions to type-specific chunk processing.
type Chunk struct {
	*ChunkOp
	KeyValue
}

// Requirements lists required backend interfaces for a type.
type Requirements struct {
	BulkIniter bool
	BulkWriter bool
	Batcher    bool
}

// Engine implementations can fulfill a variety of interfaces.  Other parts
// of DVID, most notably the data type implementations, need to know what's available.
// Data types can throw a warning at init time if the backend doesn't support required
// interfaces, or they can choose to implement multiple ways of handling data.
type Engine interface {
	KeyValueDB

	IsBatcher() bool
	IsBulkIniter() bool
	IsBulkWriter() bool

	GetConfig() dvid.Config
}

// KeyValueDB provides an interface to the simplest storage API: a key/value store.
type KeyValueDB interface {
	// Closes datastore.
	Close()

	// Get returns a value given a key.
	Get(k Key) (v []byte, err error)

	// GetRange returns a range of values spanning (kStart, kEnd) keys.
	GetRange(kStart, kEnd Key) (values []KeyValue, err error)

	// KeysInRange returns a range of keys spanning (kStart, kEnd).
	KeysInRange(kStart, kEnd Key) (keys []Key, err error)

	// ProcessRange sends a range of key/value pairs to type-specific chunk handlers.
	ProcessRange(kStart, kEnd Key, op *ChunkOp, f func(*Chunk)) (err error)

	// Put writes a value with given key.
	Put(k Key, v []byte) error

	// Put key-value pairs.  Note that it could be more efficient to use the Batcher
	// interface so you don't have to create and keep a slice of KeyValue.  Some
	// databases like leveldb will copy on batch put anyway.
	PutRange(values []KeyValue) error

	// Delete removes an entry given key.
	Delete(k Key) error
}

// Batchers allow batching operations into an atomic update or transaction.
// For example: "Atomic Updates" in http://leveldb.googlecode.com/svn/trunk/doc/index.html
type Batcher interface {
	NewBatch() Batch
}

// Batch groups operations into a transaction.
type Batch interface {
	// Commits a batch of operations.
	Commit() error

	// Delete removes from the batch a put using the given key.
	Delete(k Key)

	// Put adds to the batch a put using the given key/value.
	Put(k Key, v []byte)

	// Clear clears the contents of a write batch
	Clear()

	// Close closes a write batch
	Close()
}

// BulkIniters can employ even more aggressive optimization in loading large
// data since they can assume an uninitialized blank database.
type BulkIniter interface {
}

// BulkWriter employ some sort of optimization to efficiently write large
// amount of data.  For some key-value databases, this requires keys to
// be presorted.
type BulkWriter interface {
}
<commit_msg>Separate getter and setter functions in key-value interface<commit_after>/*
	Package storage provides a unified interface to a number of storage engines.
	Since each storage engine has different capabilities, this package defines a
	number of interfaces and the Engine interface provides a way to query
	which interfaces are implemented by a given storage engine.

	Initially we are concentrating on key-value storage engines but expect to
	expand support to graph and relational databases.

	Each storage engine must implement the following:

		NewStore(path string, create bool, options *Options) (Engine, error)
*/
package storage

import (
	"sync"

	"github.com/janelia-flyem/dvid/dvid"
)

func Shutdown() {
	ShutdownFUSE()
}

// ChunkOp is a type-specific operation with an optional WaitGroup to
// sync mapping before reduce.
type ChunkOp struct {
	Op interface{}
	Wg *sync.WaitGroup
}

// Chunk is the unit passed down channels to chunk handlers.  Chunks can be passed
// from lower-level database access functions to type-specific chunk processing.
type Chunk struct {
	*ChunkOp
	KeyValue
}

// Requirements lists required backend interfaces for a type.
type Requirements struct {
	BulkIniter bool
	BulkWriter bool
	Batcher    bool
}

// Engine implementations can fulfill a variety of interfaces.  Other parts
// of DVID, most notably the data type implementations, need to know what's available.
// Data types can throw a warning at init time if the backend doesn't support required
// interfaces, or they can choose to implement multiple ways of handling data.
type Engine interface {
	KeyValueDB

	IsBatcher() bool
	IsBulkIniter() bool
	IsBulkWriter() bool

	GetConfig() dvid.Config
}

type KeyValueGetter interface {
	// Get returns a value given a key.
	Get(k Key) (v []byte, err error)

	// GetRange returns a range of values spanning (kStart, kEnd) keys.
	GetRange(kStart, kEnd Key) (values []KeyValue, err error)

	// KeysInRange returns a range of keys spanning (kStart, kEnd).
	KeysInRange(kStart, kEnd Key) (keys []Key, err error)

	// ProcessRange sends a range of key/value pairs to type-specific chunk handlers.
	ProcessRange(kStart, kEnd Key, op *ChunkOp, f func(*Chunk)) (err error)
}

type KeyValueSetter interface {
	// Put writes a value with given key.
	Put(k Key, v []byte) error

	// Put key-value pairs.  Note that it could be more efficient to use the Batcher
	// interface so you don't have to create and keep a slice of KeyValue.  Some
	// databases like leveldb will copy on batch put anyway.
	PutRange(values []KeyValue) error

	// Delete removes an entry given key.
	Delete(k Key) error
}

// KeyValueDB provides an interface to the simplest storage API: a key/value store.
type KeyValueDB interface {
	KeyValueGetter
	KeyValueSetter

	// Closes datastore.
	Close()
}

// Batchers allow batching operations into an atomic update or transaction.
// For example: "Atomic Updates" in http://leveldb.googlecode.com/svn/trunk/doc/index.html
type Batcher interface {
	NewBatch() Batch
}

// Batch groups operations into a transaction.
type Batch interface {
	// Commits a batch of operations.
	Commit() error

	// Delete removes from the batch a put using the given key.
	Delete(k Key)

	// Put adds to the batch a put using the given key/value.
	Put(k Key, v []byte)

	// Clear clears the contents of a write batch
	Clear()

	// Close closes a write batch
	Close()
}

// BulkIniters can employ even more aggressive optimization in loading large
// data since they can assume an uninitialized blank database.
type BulkIniter interface {
}

// BulkWriter employ some sort of optimization to efficiently write large
// amount of data.  For some key-value databases, this requires keys to
// be presorted.
type BulkWriter interface {
}
<|endoftext|>


<commit_before>// Copyright 2014 mqant Author. All Rights Reserved.
//
// 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 mqtt

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/liangdas/mqant/conf"
	"github.com/liangdas/mqant/network"
	"runtime"
	"sync"
	"time"
)

// Tcp write queue
type PackQueue struct {
	conf conf.Mqtt
	// The last error in the tcp connection
	writeError error
	// Notice read the error
	fch       chan struct{}
	writelock sync.Mutex
	recover   func(pAndErr *packAndErr) (err error)
	// Pack connection
	r *bufio.Reader
	w *bufio.Writer

	conn network.Conn

	alive int

	status int
}

type packAndErr struct {
	pack *Pack
	err  error
}

// 1 is delay, 0 is no delay, 2 is just flush.
const (
	NO_DELAY = iota
	DELAY
	FLUSH

	DISCONNECTED = iota
	CONNECTED
	CLOSED
	RECONNECTING
	CONNECTING
)

type packAndType struct {
	pack *Pack
	typ  byte
}

// Init a pack queue
func NewPackQueue(conf conf.Mqtt, r *bufio.Reader, w *bufio.Writer, conn network.Conn, recover func(pAndErr *packAndErr) (err error), alive int) *PackQueue {
	if alive < 1 {
		alive = conf.ReadTimeout
	}
	alive = int(float32(alive)*1.5 + 1)
	return &PackQueue{
		conf:    conf,
		alive:   alive,
		r:       r,
		w:       w,
		conn:    conn,
		recover: recover,
		fch:     make(chan struct{}, 256),
		status:  CONNECTED,
	}
}

func (queue *PackQueue) isConnected() bool {
	return queue.status == CONNECTED
}

// Get a read pack queue
// Only call once
func (queue *PackQueue) Flusher() {
	for queue.isConnected() {
		if _, ok := <-queue.fch; !ok {
			break
		}
		queue.writelock.Lock()
		if !queue.isConnected() {
			queue.writelock.Unlock()
			break
		}
		if queue.w.Buffered() > 0 {
			if err := queue.w.Flush(); err != nil {
				queue.writelock.Unlock()
				break
			}
		}
		queue.writelock.Unlock()
	}
	//log.Info("flusher_loop Groutine will esc.")
}

// Write a pack , and get the last error
func (queue *PackQueue) WritePack(pack *Pack) (err error) {
	defer func() {
		if r := recover(); r != nil {
			buf := make([]byte, 1024)
			l := runtime.Stack(buf, false)
			errstr := string(buf[:l])
			err = fmt.Errorf("WritePack error %v", errstr)
			queue.Close(err)
		}
	}()
	queue.writelock.Lock()
	if !queue.isConnected() {
		queue.writelock.Unlock()
		return errors.New("disconnect")
	}
	queue.writelock.Unlock()
	if queue.writeError != nil {
		return queue.writeError
	}
	queue.writelock.Lock()
	defer queue.writelock.Unlock()
	err = DelayWritePack(pack, queue.w)
	queue.fch <- struct{}{}
	if err != nil {
		// Tell listener the error
		// Notice the read
		queue.Close(err)
	}
	return err
}

func (queue *PackQueue) SetAlive(alive int) error {
	if alive < 1 {
		alive = queue.conf.ReadTimeout
	}
	alive = int(float32(alive)*1.5 + 1)
	queue.alive = alive
	return nil
}

// Get a read pack queue
// Only call once
func (queue *PackQueue) ReadPackInLoop() {
	// defer recover()
	p := new(packAndErr)
loop:
	for queue.isConnected() {
		if queue.alive > 0 {
			queue.conn.SetDeadline(time.Now().Add(time.Second * time.Duration(int(float64(queue.alive)*1.5))))
		} else {
			queue.conn.SetDeadline(time.Now().Add(time.Second * 90))
		}
		p.pack, p.err = ReadPack(queue.r)
		if p.err != nil {
			queue.Close(p.err)
			break loop
		}
		err := queue.recover(p)
		if err != nil {
			queue.Close(err)
			break loop
		}
		p = new(packAndErr)
	}

	//log.Info("read_loop Groutine will esc.")
}
func (queue *PackQueue) CloseFch() {
	defer func() {
		if recover() != nil {
			// close(ch) panic occur
		}
	}()

	close(queue.fch) // panic if ch is closed
}

// Close the all of queue's channels
func (queue *PackQueue) Close(err error) error {
	queue.writeError = err
	queue.CloseFch()
	queue.status = CLOSED
	return nil
}

// Buffer
type buffer struct {
	index int
	data  []byte
}

func newBuffer(data []byte) *buffer {
	return &buffer{
		data:  data,
		index: 0,
	}
}
func (b *buffer) readString(length int) (s string, err error) {
	if (length + b.index) > len(b.data) {
		err = fmt.Errorf("Out of range error:%v", length)
		return
	}
	s = string(b.data[b.index:(length + b.index)])
	b.index += length
	return
}
func (b *buffer) readByte() (c byte, err error) {
	if (1 + b.index) > len(b.data) {
		err = fmt.Errorf("Out of range error")
		return
	}
	c = b.data[b.index]
	b.index++
	return
}
<commit_msg>=优化mqtt网关<commit_after>// Copyright 2014 mqant Author. All Rights Reserved.
//
// 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 mqtt

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/liangdas/mqant/conf"
	"github.com/liangdas/mqant/network"
	"runtime"
	"sync"
	"time"
)

// Tcp write queue
type PackQueue struct {
	conf conf.Mqtt
	// The last error in the tcp connection
	writeError error
	// Notice read the error
	fch       chan struct{}
	writelock sync.Mutex
	recover   func(pAndErr *packAndErr) (err error)
	// Pack connection
	r *bufio.Reader
	w *bufio.Writer

	conn network.Conn

	alive int

	status int
}

type packAndErr struct {
	pack *Pack
	err  error
}

// 1 is delay, 0 is no delay, 2 is just flush.
const (
	NO_DELAY = iota
	DELAY
	FLUSH

	DISCONNECTED = iota
	CONNECTED
	CLOSED
	RECONNECTING
	CONNECTING
)

type packAndType struct {
	pack *Pack
	typ  byte
}

// Init a pack queue
func NewPackQueue(conf conf.Mqtt, r *bufio.Reader, w *bufio.Writer, conn network.Conn, recover func(pAndErr *packAndErr) (err error), alive int) *PackQueue {
	if alive < 1 {
		alive = conf.ReadTimeout
	}
	alive = int(float32(alive)*1.5 + 1)
	return &PackQueue{
		conf:    conf,
		alive:   alive,
		r:       r,
		w:       w,
		conn:    conn,
		recover: recover,
		fch:     make(chan struct{}, 256),
		status:  CONNECTED,
	}
}

func (queue *PackQueue) isConnected() bool {
	return queue.status == CONNECTED
}

// Get a read pack queue
// Only call once
func (queue *PackQueue) Flusher() {
	for queue.isConnected() {
		if _, ok := <-queue.fch; !ok {
			break
		}
		queue.writelock.Lock()
		if !queue.isConnected() {
			queue.writelock.Unlock()
			break
		}
		if queue.w.Buffered() > 0 {
			if err := queue.w.Flush(); err != nil {
				queue.writelock.Unlock()
				break
			}
		}
		queue.writelock.Unlock()
	}
	//log.Info("flusher_loop Groutine will esc.")
}

// Write a pack , and get the last error
func (queue *PackQueue) WritePack(pack *Pack) (err error) {
	defer func() {
		if r := recover(); r != nil {
			buf := make([]byte, 1024)
			l := runtime.Stack(buf, false)
			errstr := string(buf[:l])
			err = fmt.Errorf("WritePack error %v", errstr)
			queue.Close(err)
		}
	}()
	queue.writelock.Lock()
	if !queue.isConnected() {
		queue.writelock.Unlock()
		return errors.New("disconnect")
	}
	if queue.writeError != nil {
		queue.writelock.Unlock()
		return queue.writeError
	}
	err = DelayWritePack(pack, queue.w)
	queue.writelock.Unlock()
	queue.fch <- struct{}{}
	if err != nil {
		// Tell listener the error
		// Notice the read
		queue.Close(err)
	}
	return err
}

func (queue *PackQueue) SetAlive(alive int) error {
	if alive < 1 {
		alive = queue.conf.ReadTimeout
	}
	alive = int(float32(alive)*1.5 + 1)
	queue.alive = alive
	return nil
}

// Get a read pack queue
// Only call once
func (queue *PackQueue) ReadPackInLoop() {
	// defer recover()
	p := new(packAndErr)
loop:
	for queue.isConnected() {
		if queue.alive > 0 {
			queue.conn.SetDeadline(time.Now().Add(time.Second * time.Duration(int(float64(queue.alive)*1.5))))
		} else {
			queue.conn.SetDeadline(time.Now().Add(time.Second * 90))
		}
		p.pack, p.err = ReadPack(queue.r)
		if p.err != nil {
			queue.Close(p.err)
			break loop
		}
		err := queue.recover(p)
		if err != nil {
			queue.Close(err)
			break loop
		}
		p = new(packAndErr)
	}

	//log.Info("read_loop Groutine will esc.")
}
func (queue *PackQueue) CloseFch() {
	defer func() {
		if recover() != nil {
			// close(ch) panic occur
		}
	}()

	close(queue.fch) // panic if ch is closed
}

// Close the all of queue's channels
func (queue *PackQueue) Close(err error) error {
	queue.writeError = err
	queue.CloseFch()
	queue.status = CLOSED
	return nil
}

// Buffer
type buffer struct {
	index int
	data  []byte
}

func newBuffer(data []byte) *buffer {
	return &buffer{
		data:  data,
		index: 0,
	}
}
func (b *buffer) readString(length int) (s string, err error) {
	if (length + b.index) > len(b.data) {
		err = fmt.Errorf("Out of range error:%v", length)
		return
	}
	s = string(b.data[b.index:(length + b.index)])
	b.index += length
	return
}
func (b *buffer) readByte() (c byte, err error) {
	if (1 + b.index) > len(b.data) {
		err = fmt.Errorf("Out of range error")
		return
	}
	c = b.data[b.index]
	b.index++
	return
}
<|endoftext|>


<commit_before>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

/* Package cover implements set operations on slices (arrays of coverage PCs). */
package cover

import (
	"sort"
)

type Cover []uint32

func (a Cover) Len() int           { return len(a) }
func (a Cover) Less(i, j int) bool { return a[i] < a[j] }
func (a Cover) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

const sent = ^uint32(0)

func Copy(cov Cover) Cover {
	return append(Cover{}, cov...)
}

func RestorePC(pc uint32, base uint32) uint64 {
	return uint64(base)<<32 + uint64(pc)
}

/* Canonicalize sorts and removes duplicates. */
func Canonicalize(cov []uint32) Cover {
	sort.Sort(Cover(cov))
	i := 0
	last := sent
	for _, pc := range cov {
		if pc != last {
			last = pc
			cov[i] = pc
			i++
		}
	}
	return Cover(cov[:i])
}

func Difference(cov0, cov1 Cover) Cover {
	return foreach(cov0, cov1, func(v0, v1 uint32) uint32 {
		if v0 < v1 {
			return v0
		}
		return sent
	})
}

func SymmetricDifference(cov0, cov1 Cover) Cover {
	return foreach(cov0, cov1, func(v0, v1 uint32) uint32 {
		if v0 < v1 {
			return v0
		}
		if v1 < v0 {
			return v1
		}
		return sent
	})
}

func Union(cov0, cov1 Cover) Cover {
	return foreach(cov0, cov1, func(v0, v1 uint32) uint32 {
		if v0 <= v1 {
			return v0
		}
		return v1
	})
}

func Intersection(cov0, cov1 Cover) Cover {
	return foreach(cov0, cov1, func(v0, v1 uint32) uint32 {
		if v0 == v1 {
			return v0
		}
		return sent
	})
}

func foreach(cov0, cov1 Cover, f func(uint32, uint32) uint32) Cover {
	var res []uint32
	for i0, i1 := 0, 0; i0 < len(cov0) || i1 < len(cov1); {
		v0, v1 := sent, sent
		if i0 < len(cov0) {
			v0 = cov0[i0]
		}
		if i1 < len(cov1) {
			v1 = cov1[i1]
		}
		if v0 <= v1 {
			i0++
		}
		if v1 <= v0 {
			i1++
		}
		if v := f(v0, v1); v != sent {
			res = append(res, v)
		}
	}
	return res
}

// Minimize returns a minimal set of inputs that give the same coverage as the full corpus.
func Minimize(corpus []Cover) []int {
	inputs := make([]*minInput, len(corpus))
	for i, cov := range corpus {
		inputs[i] = &minInput{
			idx: i,
			cov: cov,
		}
	}
	sort.Sort(minInputArray(inputs))
	var min []int
	covered := make(map[uint32]struct{})
	for _, inp := range inputs {
		hit := false
		for _, pc := range inp.cov {
			if !hit {
				if _, ok := covered[pc]; !ok {
					hit = true
					min = append(min, inp.idx)
				}
			}
			if hit {
				covered[pc] = struct{}{}
			}
		}
	}
	return min
}

type minInput struct {
	idx int
	cov Cover
}

type minInputArray []*minInput

// Inputs with larger coverage come first.
func (a minInputArray) Len() int           { return len(a) }
func (a minInputArray) Less(i, j int) bool { return len(a[i].cov) > len(a[j].cov) }
func (a minInputArray) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
<commit_msg>cover: fix comment style<commit_after>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

/* Package cover implements set operations on slices (arrays of coverage PCs). */
package cover

import (
	"sort"
)

type Cover []uint32

func (a Cover) Len() int           { return len(a) }
func (a Cover) Less(i, j int) bool { return a[i] < a[j] }
func (a Cover) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

const sent = ^uint32(0)

func Copy(cov Cover) Cover {
	return append(Cover{}, cov...)
}

func RestorePC(pc uint32, base uint32) uint64 {
	return uint64(base)<<32 + uint64(pc)
}

// Canonicalize sorts and removes duplicates.
func Canonicalize(cov []uint32) Cover {
	sort.Sort(Cover(cov))
	i := 0
	last := sent
	for _, pc := range cov {
		if pc != last {
			last = pc
			cov[i] = pc
			i++
		}
	}
	return Cover(cov[:i])
}

func Difference(cov0, cov1 Cover) Cover {
	return foreach(cov0, cov1, func(v0, v1 uint32) uint32 {
		if v0 < v1 {
			return v0
		}
		return sent
	})
}

func SymmetricDifference(cov0, cov1 Cover) Cover {
	return foreach(cov0, cov1, func(v0, v1 uint32) uint32 {
		if v0 < v1 {
			return v0
		}
		if v1 < v0 {
			return v1
		}
		return sent
	})
}

func Union(cov0, cov1 Cover) Cover {
	return foreach(cov0, cov1, func(v0, v1 uint32) uint32 {
		if v0 <= v1 {
			return v0
		}
		return v1
	})
}

func Intersection(cov0, cov1 Cover) Cover {
	return foreach(cov0, cov1, func(v0, v1 uint32) uint32 {
		if v0 == v1 {
			return v0
		}
		return sent
	})
}

func foreach(cov0, cov1 Cover, f func(uint32, uint32) uint32) Cover {
	var res []uint32
	for i0, i1 := 0, 0; i0 < len(cov0) || i1 < len(cov1); {
		v0, v1 := sent, sent
		if i0 < len(cov0) {
			v0 = cov0[i0]
		}
		if i1 < len(cov1) {
			v1 = cov1[i1]
		}
		if v0 <= v1 {
			i0++
		}
		if v1 <= v0 {
			i1++
		}
		if v := f(v0, v1); v != sent {
			res = append(res, v)
		}
	}
	return res
}

// Minimize returns a minimal set of inputs that give the same coverage as the full corpus.
func Minimize(corpus []Cover) []int {
	inputs := make([]*minInput, len(corpus))
	for i, cov := range corpus {
		inputs[i] = &minInput{
			idx: i,
			cov: cov,
		}
	}
	sort.Sort(minInputArray(inputs))
	var min []int
	covered := make(map[uint32]struct{})
	for _, inp := range inputs {
		hit := false
		for _, pc := range inp.cov {
			if !hit {
				if _, ok := covered[pc]; !ok {
					hit = true
					min = append(min, inp.idx)
				}
			}
			if hit {
				covered[pc] = struct{}{}
			}
		}
	}
	return min
}

type minInput struct {
	idx int
	cov Cover
}

type minInputArray []*minInput

// Inputs with larger coverage come first.
func (a minInputArray) Len() int           { return len(a) }
func (a minInputArray) Less(i, j int) bool { return len(a[i].cov) > len(a[j].cov) }
func (a minInputArray) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
<|endoftext|>