<commit_before>/*
 * Minimal object storage library (C) 2015 Minio, Inc.
 *
 * 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 objectstorage

import (
	"errors"
	"io"
	"net/http"
	"runtime"
	"sort"
	"strings"
	"time"
)

// API - object storage API interface
type API interface {
	// Bucket Read/Write/Stat operations
	BucketAPI

	// Object Read/Write/Stat operations
	ObjectAPI
}

// BucketAPI - bucket specific Read/Write/Stat interface
type BucketAPI interface {
	CreateBucket(bucket, acl, location string) error
	SetBucketACL(bucket, acl string) error
	StatBucket(bucket string) error
	DeleteBucket(bucket string) error

	ListObjects(bucket, prefix string, recursive bool) <-chan ObjectOnChannel
	ListBuckets() <-chan BucketOnChannel
}

// ObjectAPI - object specific Read/Write/Stat interface
type ObjectAPI interface {
	GetObject(bucket, object string, offset, length uint64) (io.ReadCloser, *ObjectMetadata, error)
	CreateObject(bucket, object string, size uint64, data io.Reader) (string, error)
	StatObject(bucket, object string) (*ObjectMetadata, error)
	DeleteObject(bucket, object string) error
}

// BucketOnChannel - bucket metadata over read channel
type BucketOnChannel struct {
	Data *BucketMetadata
	Err  error
}

// ObjectOnChannel - object metadata over read channel
type ObjectOnChannel struct {
	Data *ObjectMetadata
	Err  error
}

// BucketMetadata container for bucket metadata
type BucketMetadata struct {
	// The name of the bucket.
	Name string
	// Date the bucket was created.
	CreationDate time.Time
}

// ObjectMetadata container for object metadata
type ObjectMetadata struct {
	ETag         string
	Key          string
	LastModified time.Time
	Size         int64

	Owner struct {
		DisplayName string
		ID          string
	}

	// The class of storage used to store the object.
	StorageClass string
}

// Regions s3 region map used by bucket location constraint
var Regions = map[string]string{
	"us-gov-west-1":  "https://s3-fips-us-gov-west-1.amazonaws.com",
	"us-east-1":      "https://s3.amazonaws.com",
	"us-west-1":      "https://s3-us-west-1.amazonaws.com",
	"us-west-2":      "https://s3-us-west-2.amazonaws.com",
	"eu-west-1":      "https://s3-eu-west-1.amazonaws.com",
	"eu-central-1":   "https://s3-eu-central-1.amazonaws.com",
	"ap-southeast-1": "https://s3-ap-southeast-1.amazonaws.com",
	"ap-southeast-2": "https://s3-ap-southeast-2.amazonaws.com",
	"ap-northeast-1": "https://s3-ap-northeast-1.amazonaws.com",
	"sa-east-1":      "https://s3-sa-east-1.amazonaws.com",
	"cn-north-1":     "https://s3.cn-north-1.amazonaws.com.cn",
}

// getEndpoint fetches an endpoint based on region through the S3 Regions map
func getEndpoint(region string) string {
	return Regions[region]
}

type api struct {
	*lowLevelAPI
}

// Config - main configuration struct used by all to set endpoint, credentials, and other options for requests.
type Config struct {
	// Standard options
	AccessKeyID     string
	SecretAccessKey string
	Region          string
	Endpoint        string

	// Advanced options
	AcceptType string            // specify this to get server response in non XML style if server supports it
	UserAgent  string            // user override useful when objectstorage-go is used with in your application
	Transport  http.RoundTripper // custom transport usually for debugging, by default its nil
}

// MustGetEndpoint makes sure that a valid endpoint is provided all the time, even with false regions it will fall
// back to default, for no regions specified it chooses to default to "milkyway" and use endpoint as is
func (c *Config) MustGetEndpoint() string {
	// for custom domains, there are no regions default to 'milkyway'
	if strings.TrimSpace(c.Region) == "" && strings.TrimSpace(c.Endpoint) != "" {
		c.Region = "milkyway"
		return c.Endpoint
	}
	// if valid region provided override user provided endpoint
	if endpoint := getEndpoint(strings.TrimSpace(c.Region)); endpoint != "" {
		c.Endpoint = endpoint
		return c.Endpoint
	}
	// fall back if region is set and not found, go through US-standard
	c.Region = "us-east-1"
	return getEndpoint(c.Region)
}

// Global constants
const (
	LibraryName    = "objectstorage-go/"
	LibraryVersion = "0.1"
)

// New - instantiate a new minio api client
func New(config *Config) API {
	// if not UserAgent provided set it to default
	if strings.TrimSpace(config.UserAgent) == "" {
		config.UserAgent = LibraryName + " (" + LibraryVersion + "; " + runtime.GOOS + "; " + runtime.GOARCH + ")"
	}
	return &api{&lowLevelAPI{config}}
}

/// Object operations

// GetObject retrieve object
//
// Additionally it also takes range arguments to download the specified range bytes of an object.
// For more information about the HTTP Range header, go to http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.35.
func (a *api) GetObject(bucket, object string, offset, length uint64) (io.ReadCloser, *ObjectMetadata, error) {
	// get the the object
	// NOTE : returned md5sum could be the md5sum of the partial object itself
	// not the whole object depending on if offset range was requested or not
	body, objectMetadata, err := a.getObject(bucket, object, offset, length)
	if err != nil {
		return nil, nil, err
	}
	return body, objectMetadata, nil
}

// completedParts is a wrapper to make parts sortable by their part number
// multi part completion requires list of multi parts to be sorted
type completedParts []*completePart

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

// DefaultPartSize - default size per object after which PutObject becomes multipart
// one can change this value during a library import
var DefaultPartSize uint64 = 1024 * 1024 * 5

// CreateObject create an object in a bucket
//
// You must have WRITE permissions on a bucket to create an object
//
// This version of CreateObject automatically does multipart for more than 5MB worth of data
func (a *api) CreateObject(bucket, object string, size uint64, data io.Reader) (string, error) {
	switch {
	case size < DefaultPartSize:
		// Single Part use case, use PutObject directly
		for part := range MultiPart(data, DefaultPartSize) {
			if part.Err != nil {
				return "", part.Err
			}
			metadata, err := a.putObject(bucket, object, part.Len, part.Data)
			if err != nil {
				return "", err
			}
			return metadata.ETag, nil
		}
	default:
		initiateMultipartUploadResult, err := a.initiateMultipartUpload(bucket, object)
		if err != nil {
			return "", err
		}
		uploadID := initiateMultipartUploadResult.UploadID
		completeMultipartUpload := new(completeMultipartUpload)
		for part := range MultiPart(data, DefaultPartSize) {
			if part.Err != nil {
				return "", part.Err
			}
			completePart, err := a.uploadPart(bucket, object, uploadID, part.Num, part.Len, part.Data)
			if err != nil {
				return "", a.abortMultipartUpload(bucket, object, uploadID)
			}
			completeMultipartUpload.Part = append(completeMultipartUpload.Part, completePart)
		}
		sort.Sort(completedParts(completeMultipartUpload.Part))
		completeMultipartUploadResult, err := a.completeMultipartUpload(bucket, object, uploadID, completeMultipartUpload)
		if err != nil {
			return "", a.abortMultipartUpload(bucket, object, uploadID)
		}
		return completeMultipartUploadResult.ETag, nil
	}
	return "", errors.New("Unexpected control flow")
}

// StatObject verify if object exists and you have permission to access it
func (a *api) StatObject(bucket, object string) (*ObjectMetadata, error) {
	return a.headObject(bucket, object)
}

// DeleteObject remove the object from a bucket
func (a *api) DeleteObject(bucket, object string) error {
	return a.deleteObject(bucket, object)
}

/// Bucket operations

// CreateBucket create a new bucket
//
// optional arguments are acl and location - by default all buckets are created
// with ``private`` acl and location set to US Standard if one wishes to set
// different ACLs and Location one can set them properly.
//
// ACL valid values
// ------------------
// private - owner gets full access [DEFAULT]
// public-read - owner gets full access, others get read access
// public-read-write - owner gets full access, others get full access too
// ------------------
//
// Location valid values
// ------------------
// [ us-west-1 | us-west-2 | eu-west-1 | eu-central-1 | ap-southeast-1 | ap-northeast-1 | ap-southeast-2 | sa-east-1 ]
// Default - US standard
func (a *api) CreateBucket(bucket, acl, location string) error {
	return a.putBucket(bucket, acl, location)
}

// SetBucketACL set the permissions on an existing bucket using access control lists (ACL)
//
// Currently supported are:
// ------------------
// private - owner gets full access
// public-read - owner gets full access, others get read access
// public-read-write - owner gets full access, others get full access too
// ------------------
func (a *api) SetBucketACL(bucket, acl string) error {
	return a.putBucketACL(bucket, acl)
}

// StatBucket verify if bucket exists and you have permission to access it
func (a *api) StatBucket(bucket string) error {
	return a.headBucket(bucket)
}

// DeleteBucket deletes the bucket named in the URI
// NOTE: -
//  All objects (including all object versions and delete markers)
//  in the bucket must be deleted before successfully attempting this request
func (a *api) DeleteBucket(bucket string) error {
	return a.deleteBucket(bucket)
}

// listObjectsInRoutine is an internal goroutine function called for listing objects
// This function feeds data into channel
func (a *api) listObjectsInRoutine(bucket, prefix string, recursive bool, ch chan ObjectOnChannel) {
	defer close(ch)
	switch {
	case recursive == true:
		listBucketResult, err := a.listObjects(bucket, 1000, "", prefix, "")
		if err != nil {
			ch <- ObjectOnChannel{
				Data: nil,
				Err:  err,
			}
			return
		}
		for _, object := range listBucketResult.Contents {
			ch <- ObjectOnChannel{
				Data: object,
				Err:  nil,
			}
		}
		for {
			if !listBucketResult.IsTruncated {
				break
			}
			listBucketResult, err = a.listObjects(bucket, 1000, listBucketResult.Marker, prefix, "")
			if err != nil {
				ch <- ObjectOnChannel{
					Data: nil,
					Err:  err,
				}
				return
			}
			for _, object := range listBucketResult.Contents {
				ch <- ObjectOnChannel{
					Data: object,
					Err:  nil,
				}
				listBucketResult.Marker = object.Key
			}
		}
	default:
		listBucketResult, err := a.listObjects(bucket, 1000, "", prefix, "/")
		if err != nil {
			ch <- ObjectOnChannel{
				Data: nil,
				Err:  err,
			}
			return
		}
		for _, object := range listBucketResult.Contents {
			ch <- ObjectOnChannel{
				Data: object,
				Err:  nil,
			}
		}
		for _, prefix := range listBucketResult.CommonPrefixes {
			object := new(ObjectMetadata)
			object.Key = prefix.Prefix
			object.Size = 0
			ch <- ObjectOnChannel{
				Data: object,
				Err:  nil,
			}
		}
	}
}

// ListObjects - (List Objects) - List some objects or all recursively
//
// ListObjects is a channel based API implemented to facilitate ease of usage of S3 API ListObjects()
// by automatically recursively traversing all objects on a given bucket if specified.
//
// Your input paramters are just bucket, prefix and recursive
//
// If you enable recursive as 'true' this function will return back all the objects in a given bucket
//
//  eg:-
//         api := objectstorage.New(....)
//         for message := range api.ListObjects("mytestbucket", "starthere", true) {
//                 fmt.Println(message.Data)
//         }
//
func (a *api) ListObjects(bucket string, prefix string, recursive bool) <-chan ObjectOnChannel {
	ch := make(chan ObjectOnChannel)
	go a.listObjectsInRoutine(bucket, prefix, recursive, ch)
	return ch
}

// listBucketsInRoutine is an internal go routine function called for listing buckets
// This function feeds data into channel
func (a *api) listBucketsInRoutine(ch chan BucketOnChannel) {
	defer close(ch)
	listAllMyBucketListResults, err := a.listBuckets()
	if err != nil {
		ch <- BucketOnChannel{
			Data: nil,
			Err:  err,
		}
		return
	}
	for _, bucket := range listAllMyBucketListResults.Buckets.Bucket {
		ch <- BucketOnChannel{
			Data: bucket,
			Err:  nil,
		}
	}

}

// ListBuckets list of all buckets owned by the authenticated sender of the request
//
// NOTE:
//     This call requires explicit authentication, no anonymous
//     requests are allowed for listing buckets
//
//  eg:-
//         api := objectstorage.New(....)
//         for message := range api.ListBuckets() {
//                 fmt.Println(message.Data)
//         }
//
func (a *api) ListBuckets() <-chan BucketOnChannel {
	ch := make(chan BucketOnChannel)
	go a.listBucketsInRoutine(ch)
	return ch
}
<commit_msg>MustGetEndpoint should return valid endpoint for custom domains<commit_after>/*
 * Minimal object storage library (C) 2015 Minio, Inc.
 *
 * 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 objectstorage

import (
	"errors"
	"io"
	"net/http"
	"runtime"
	"sort"
	"strings"
	"time"
)

// API - object storage API interface
type API interface {
	// Bucket Read/Write/Stat operations
	BucketAPI

	// Object Read/Write/Stat operations
	ObjectAPI
}

// BucketAPI - bucket specific Read/Write/Stat interface
type BucketAPI interface {
	CreateBucket(bucket, acl, location string) error
	SetBucketACL(bucket, acl string) error
	StatBucket(bucket string) error
	DeleteBucket(bucket string) error

	ListObjects(bucket, prefix string, recursive bool) <-chan ObjectOnChannel
	ListBuckets() <-chan BucketOnChannel
}

// ObjectAPI - object specific Read/Write/Stat interface
type ObjectAPI interface {
	GetObject(bucket, object string, offset, length uint64) (io.ReadCloser, *ObjectMetadata, error)
	CreateObject(bucket, object string, size uint64, data io.Reader) (string, error)
	StatObject(bucket, object string) (*ObjectMetadata, error)
	DeleteObject(bucket, object string) error
}

// BucketOnChannel - bucket metadata over read channel
type BucketOnChannel struct {
	Data *BucketMetadata
	Err  error
}

// ObjectOnChannel - object metadata over read channel
type ObjectOnChannel struct {
	Data *ObjectMetadata
	Err  error
}

// BucketMetadata container for bucket metadata
type BucketMetadata struct {
	// The name of the bucket.
	Name string
	// Date the bucket was created.
	CreationDate time.Time
}

// ObjectMetadata container for object metadata
type ObjectMetadata struct {
	ETag         string
	Key          string
	LastModified time.Time
	Size         int64

	Owner struct {
		DisplayName string
		ID          string
	}

	// The class of storage used to store the object.
	StorageClass string
}

// Regions s3 region map used by bucket location constraint
var Regions = map[string]string{
	"us-gov-west-1":  "https://s3-fips-us-gov-west-1.amazonaws.com",
	"us-east-1":      "https://s3.amazonaws.com",
	"us-west-1":      "https://s3-us-west-1.amazonaws.com",
	"us-west-2":      "https://s3-us-west-2.amazonaws.com",
	"eu-west-1":      "https://s3-eu-west-1.amazonaws.com",
	"eu-central-1":   "https://s3-eu-central-1.amazonaws.com",
	"ap-southeast-1": "https://s3-ap-southeast-1.amazonaws.com",
	"ap-southeast-2": "https://s3-ap-southeast-2.amazonaws.com",
	"ap-northeast-1": "https://s3-ap-northeast-1.amazonaws.com",
	"sa-east-1":      "https://s3-sa-east-1.amazonaws.com",
	"cn-north-1":     "https://s3.cn-north-1.amazonaws.com.cn",
}

// getEndpoint fetches an endpoint based on region through the S3 Regions map
func getEndpoint(region string) string {
	return Regions[region]
}

type api struct {
	*lowLevelAPI
}

// Config - main configuration struct used by all to set endpoint, credentials, and other options for requests.
type Config struct {
	// Standard options
	AccessKeyID     string
	SecretAccessKey string
	Region          string
	Endpoint        string

	// Advanced options
	AcceptType string            // specify this to get server response in non XML style if server supports it
	UserAgent  string            // user override useful when objectstorage-go is used with in your application
	Transport  http.RoundTripper // custom transport usually for debugging, by default its nil
}

// MustGetEndpoint makes sure that a valid endpoint is provided all the time, even with false regions it will fall
// back to default, for no regions specified it chooses to default to "milkyway" and use endpoint as is
func (c *Config) MustGetEndpoint() string {
	// if valid region provided override user provided endpoint
	if strings.TrimSpace(c.Region) != "" {
		if endpoint := getEndpoint(strings.TrimSpace(c.Region)); endpoint != "" {
			c.Endpoint = endpoint
			return c.Endpoint
		}
		// fall back if region is set and not found, go through US-standard
		c.Region = "us-east-1"
		return getEndpoint(c.Region)
	}
	if strings.TrimSpace(c.Endpoint) != "" {
		if strings.Contains(strings.TrimSpace(c.Endpoint), "s3.amazonaws.com") {
			c.Region = "us-east-1"
			return getEndpoint(c.Region)
		}
		// for custom domains, there are no regions default to 'milkyway'
		c.Region = "milkyway"
		return c.Endpoint
	}
	// if not endpoint or region sepcified default to us-east-1
	c.Region = "us-east-1"
	return getEndpoint(c.Region)
}

// Global constants
const (
	LibraryName    = "objectstorage-go/"
	LibraryVersion = "0.1"
)

// New - instantiate a new minio api client
func New(config *Config) API {
	// if not UserAgent provided set it to default
	if strings.TrimSpace(config.UserAgent) == "" {
		config.UserAgent = LibraryName + " (" + LibraryVersion + "; " + runtime.GOOS + "; " + runtime.GOARCH + ")"
	}
	return &api{&lowLevelAPI{config}}
}

/// Object operations

// GetObject retrieve object
//
// Additionally it also takes range arguments to download the specified range bytes of an object.
// For more information about the HTTP Range header, go to http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.35.
func (a *api) GetObject(bucket, object string, offset, length uint64) (io.ReadCloser, *ObjectMetadata, error) {
	// get the the object
	// NOTE : returned md5sum could be the md5sum of the partial object itself
	// not the whole object depending on if offset range was requested or not
	body, objectMetadata, err := a.getObject(bucket, object, offset, length)
	if err != nil {
		return nil, nil, err
	}
	return body, objectMetadata, nil
}

// completedParts is a wrapper to make parts sortable by their part number
// multi part completion requires list of multi parts to be sorted
type completedParts []*completePart

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

// DefaultPartSize - default size per object after which PutObject becomes multipart
// one can change this value during a library import
var DefaultPartSize uint64 = 1024 * 1024 * 5

// CreateObject create an object in a bucket
//
// You must have WRITE permissions on a bucket to create an object
//
// This version of CreateObject automatically does multipart for more than 5MB worth of data
func (a *api) CreateObject(bucket, object string, size uint64, data io.Reader) (string, error) {
	switch {
	case size < DefaultPartSize:
		// Single Part use case, use PutObject directly
		for part := range MultiPart(data, DefaultPartSize) {
			if part.Err != nil {
				return "", part.Err
			}
			metadata, err := a.putObject(bucket, object, part.Len, part.Data)
			if err != nil {
				return "", err
			}
			return metadata.ETag, nil
		}
	default:
		initiateMultipartUploadResult, err := a.initiateMultipartUpload(bucket, object)
		if err != nil {
			return "", err
		}
		uploadID := initiateMultipartUploadResult.UploadID
		completeMultipartUpload := new(completeMultipartUpload)
		for part := range MultiPart(data, DefaultPartSize) {
			if part.Err != nil {
				return "", part.Err
			}
			completePart, err := a.uploadPart(bucket, object, uploadID, part.Num, part.Len, part.Data)
			if err != nil {
				return "", a.abortMultipartUpload(bucket, object, uploadID)
			}
			completeMultipartUpload.Part = append(completeMultipartUpload.Part, completePart)
		}
		sort.Sort(completedParts(completeMultipartUpload.Part))
		completeMultipartUploadResult, err := a.completeMultipartUpload(bucket, object, uploadID, completeMultipartUpload)
		if err != nil {
			return "", a.abortMultipartUpload(bucket, object, uploadID)
		}
		return completeMultipartUploadResult.ETag, nil
	}
	return "", errors.New("Unexpected control flow")
}

// StatObject verify if object exists and you have permission to access it
func (a *api) StatObject(bucket, object string) (*ObjectMetadata, error) {
	return a.headObject(bucket, object)
}

// DeleteObject remove the object from a bucket
func (a *api) DeleteObject(bucket, object string) error {
	return a.deleteObject(bucket, object)
}

/// Bucket operations

// CreateBucket create a new bucket
//
// optional arguments are acl and location - by default all buckets are created
// with ``private`` acl and location set to US Standard if one wishes to set
// different ACLs and Location one can set them properly.
//
// ACL valid values
// ------------------
// private - owner gets full access [DEFAULT]
// public-read - owner gets full access, others get read access
// public-read-write - owner gets full access, others get full access too
// ------------------
//
// Location valid values
// ------------------
// [ us-west-1 | us-west-2 | eu-west-1 | eu-central-1 | ap-southeast-1 | ap-northeast-1 | ap-southeast-2 | sa-east-1 ]
// Default - US standard
func (a *api) CreateBucket(bucket, acl, location string) error {
	return a.putBucket(bucket, acl, location)
}

// SetBucketACL set the permissions on an existing bucket using access control lists (ACL)
//
// Currently supported are:
// ------------------
// private - owner gets full access
// public-read - owner gets full access, others get read access
// public-read-write - owner gets full access, others get full access too
// ------------------
func (a *api) SetBucketACL(bucket, acl string) error {
	return a.putBucketACL(bucket, acl)
}

// StatBucket verify if bucket exists and you have permission to access it
func (a *api) StatBucket(bucket string) error {
	return a.headBucket(bucket)
}

// DeleteBucket deletes the bucket named in the URI
// NOTE: -
//  All objects (including all object versions and delete markers)
//  in the bucket must be deleted before successfully attempting this request
func (a *api) DeleteBucket(bucket string) error {
	return a.deleteBucket(bucket)
}

// listObjectsInRoutine is an internal goroutine function called for listing objects
// This function feeds data into channel
func (a *api) listObjectsInRoutine(bucket, prefix string, recursive bool, ch chan ObjectOnChannel) {
	defer close(ch)
	switch {
	case recursive == true:
		listBucketResult, err := a.listObjects(bucket, 1000, "", prefix, "")
		if err != nil {
			ch <- ObjectOnChannel{
				Data: nil,
				Err:  err,
			}
			return
		}
		for _, object := range listBucketResult.Contents {
			ch <- ObjectOnChannel{
				Data: object,
				Err:  nil,
			}
		}
		for {
			if !listBucketResult.IsTruncated {
				break
			}
			listBucketResult, err = a.listObjects(bucket, 1000, listBucketResult.Marker, prefix, "")
			if err != nil {
				ch <- ObjectOnChannel{
					Data: nil,
					Err:  err,
				}
				return
			}
			for _, object := range listBucketResult.Contents {
				ch <- ObjectOnChannel{
					Data: object,
					Err:  nil,
				}
				listBucketResult.Marker = object.Key
			}
		}
	default:
		listBucketResult, err := a.listObjects(bucket, 1000, "", prefix, "/")
		if err != nil {
			ch <- ObjectOnChannel{
				Data: nil,
				Err:  err,
			}
			return
		}
		for _, object := range listBucketResult.Contents {
			ch <- ObjectOnChannel{
				Data: object,
				Err:  nil,
			}
		}
		for _, prefix := range listBucketResult.CommonPrefixes {
			object := new(ObjectMetadata)
			object.Key = prefix.Prefix
			object.Size = 0
			ch <- ObjectOnChannel{
				Data: object,
				Err:  nil,
			}
		}
	}
}

// ListObjects - (List Objects) - List some objects or all recursively
//
// ListObjects is a channel based API implemented to facilitate ease of usage of S3 API ListObjects()
// by automatically recursively traversing all objects on a given bucket if specified.
//
// Your input paramters are just bucket, prefix and recursive
//
// If you enable recursive as 'true' this function will return back all the objects in a given bucket
//
//  eg:-
//         api := objectstorage.New(....)
//         for message := range api.ListObjects("mytestbucket", "starthere", true) {
//                 fmt.Println(message.Data)
//         }
//
func (a *api) ListObjects(bucket string, prefix string, recursive bool) <-chan ObjectOnChannel {
	ch := make(chan ObjectOnChannel)
	go a.listObjectsInRoutine(bucket, prefix, recursive, ch)
	return ch
}

// listBucketsInRoutine is an internal go routine function called for listing buckets
// This function feeds data into channel
func (a *api) listBucketsInRoutine(ch chan BucketOnChannel) {
	defer close(ch)
	listAllMyBucketListResults, err := a.listBuckets()
	if err != nil {
		ch <- BucketOnChannel{
			Data: nil,
			Err:  err,
		}
		return
	}
	for _, bucket := range listAllMyBucketListResults.Buckets.Bucket {
		ch <- BucketOnChannel{
			Data: bucket,
			Err:  nil,
		}
	}

}

// ListBuckets list of all buckets owned by the authenticated sender of the request
//
// NOTE:
//     This call requires explicit authentication, no anonymous
//     requests are allowed for listing buckets
//
//  eg:-
//         api := objectstorage.New(....)
//         for message := range api.ListBuckets() {
//                 fmt.Println(message.Data)
//         }
//
func (a *api) ListBuckets() <-chan BucketOnChannel {
	ch := make(chan BucketOnChannel)
	go a.listBucketsInRoutine(ch)
	return ch
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"gopkg.in/unrolled/render.v1"
	"log"
	"net/http"
	"os"
	"strings"
)

var (
	renderer    = render.New(render.Options{})
	fleetClient = NewClientCLIWithPeer("http://192.168.81.101:4001")
	tempDir     = "./tmp"
)

func main() {
	r := mux.NewRouter().StrictSlash(false)

	api := r.PathPrefix("/api/v1").Subrouter()

	// machines collection
	machines := api.Path("/machines").Subrouter()
	machines.Methods("GET").HandlerFunc(machineAllHandler)

	// Units collection
	units := api.Path("/units").Subrouter()
	units.Methods("GET").HandlerFunc(statusAllHandler)
	units.Methods("POST").HandlerFunc(submitUnitHandler)

	// Units singular
	unit := api.PathPrefix("/units/{id}").Subrouter()
	unit.Methods("GET").HandlerFunc(statusHandler)
	unit.Methods("DELETE").HandlerFunc(destroyHandler)
	unit.Path("/start").Methods("POST").HandlerFunc(startHandler)
	unit.Path("/stop").Methods("POST").HandlerFunc(stopHandler)
	unit.Path("/load").Methods("POST").HandlerFunc(loadHandler)

	// websocket
	r.Path("/ws/journal/{id}").HandlerFunc(wsHandler)

	n := negroni.New()
	n.UseHandler(r)

	n.Run(":3000")
}

func destroyHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	log.Printf("destroy %s unit", key)
	if err := fleetClient.Destroy(key); err != nil {
		// log.Printf("unit destroy error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func startHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	log.Printf("start %s unit", key)
	if err := fleetClient.Start(key); err != nil {
		// log.Printf("unit start error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func stopHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	log.Printf("stop %s unit", key)
	if err := fleetClient.Stop(key); err != nil {
		// log.Printf("unit stop error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func loadHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	log.Printf("load %s unit", key)
	if err := fleetClient.Load(key); err != nil {
		// log.Printf("unit load error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func submitUnitHandler(w http.ResponseWriter, req *http.Request) {
	name := req.FormValue("name")
	service := req.FormValue("service")

	if _, err := os.Stat(tempDir); os.IsNotExist(err) {
		os.Mkdir(tempDir, 0755)
	}

	serviceFile := fmt.Sprintf("%s/%s", tempDir, name)
	lines := strings.Split(string(service), "\\n")

	fo, err := os.Create(serviceFile)
	if err != nil {
		log.Printf("Open file errpr: %s", err)
		renderer.JSON(w, http.StatusBadRequest, err)
		return
	}

	defer func() {
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}()

	for _, str := range lines {
		fmt.Fprintln(fo, str)
	}

	err = fleetClient.Submit(name, serviceFile)
	if err != nil {
		// log.Printf("Fleet submit error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func machineAllHandler(w http.ResponseWriter, req *http.Request) {
	status, _ := fleetClient.MachineAll()
	renderer.JSON(w, http.StatusOK, status)
}

func statusAllHandler(w http.ResponseWriter, req *http.Request) {
	status, _ := fleetClient.StatusAll()
	renderer.JSON(w, http.StatusOK, status)
}

func statusHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	status, _ := fleetClient.StatusUnit(key)
	renderer.JSON(w, http.StatusOK, status)
}

// websocket handler

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// write journal message
	key := mux.Vars(r)["id"]
	output, _ := fleetClient.JournalF(key)
	for line := range output {
		conn.WriteMessage(websocket.TextMessage, []byte(line))
	}
}
<commit_msg>service file upload<commit_after>package main

import (
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"gopkg.in/unrolled/render.v1"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
)

var (
	renderer    = render.New(render.Options{})
	fleetClient = NewClientCLIWithPeer("http://192.168.81.101:4001")
	tempDir     = "./tmp"
)

func main() {
	r := mux.NewRouter().StrictSlash(false)

	api := r.PathPrefix("/api/v1").Subrouter()

	// machines collection
	machines := api.Path("/machines").Subrouter()
	machines.Methods("GET").HandlerFunc(machineAllHandler)

	// Units collection
	units := api.Path("/units").Subrouter()
	units.Methods("GET").HandlerFunc(statusAllHandler)
	units.Methods("POST").HandlerFunc(submitUnitHandler)
	units.Path("upload").Methods("POST").HandlerFunc(uploadUnitHandler)

	// Units singular
	unit := api.PathPrefix("/units/{id}").Subrouter()
	unit.Methods("GET").HandlerFunc(statusHandler)
	unit.Methods("DELETE").HandlerFunc(destroyHandler)
	unit.Path("/start").Methods("POST").HandlerFunc(startHandler)
	unit.Path("/stop").Methods("POST").HandlerFunc(stopHandler)
	unit.Path("/load").Methods("POST").HandlerFunc(loadHandler)

	// websocket
	r.Path("/ws/journal/{id}").HandlerFunc(wsHandler)

	n := negroni.New()
	n.UseHandler(r)

	n.Run(":3000")
}

func destroyHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	log.Printf("destroy %s unit", key)
	if err := fleetClient.Destroy(key); err != nil {
		// log.Printf("unit destroy error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func startHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	log.Printf("start %s unit", key)
	if err := fleetClient.Start(key); err != nil {
		// log.Printf("unit start error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func stopHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	log.Printf("stop %s unit", key)
	if err := fleetClient.Stop(key); err != nil {
		// log.Printf("unit stop error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func loadHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	log.Printf("load %s unit", key)
	if err := fleetClient.Load(key); err != nil {
		// log.Printf("unit load error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func uploadUnitHandler(w http.ResponseWriter, req *http.Request) {
	file, header, err := req.FormFile("file")
	defer file.Close()

	serviceFile := fmt.Sprintf("%s/%s", tempDir, header.Filename)
	out, err := os.Create(serviceFile)
	if err != nil {
		log.Printf("Open file errpr: %s", err)
		renderer.JSON(w, http.StatusBadRequest, err)
		return
	}
	defer out.Close()

	// write the content from POST to the file
	_, err = io.Copy(out, file)
	if err != nil {
		fmt.Fprintln(w, err)
	}

	err = fleetClient.Submit(header.Filename, serviceFile)
	if err != nil {
		// log.Printf("Fleet submit error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})

}

func submitUnitHandler(w http.ResponseWriter, req *http.Request) {
	name := req.FormValue("name")
	service := req.FormValue("service")

	if _, err := os.Stat(tempDir); os.IsNotExist(err) {
		os.Mkdir(tempDir, 0755)
	}

	serviceFile := fmt.Sprintf("%s/%s", tempDir, name)
	lines := strings.Split(string(service), "\\n")

	fo, err := os.Create(serviceFile)
	if err != nil {
		log.Printf("Open file errpr: %s", err)
		renderer.JSON(w, http.StatusBadRequest, err)
		return
	}

	defer func() {
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}()

	for _, str := range lines {
		fmt.Fprintln(fo, str)
	}

	err = fleetClient.Submit(name, serviceFile)
	if err != nil {
		// log.Printf("Fleet submit error: %s", err)
		// renderer.JSON(w, http.StatusBadRequest, err)
		// return
	}
	renderer.JSON(w, http.StatusOK, map[string]string{"result": "success"})
}

func machineAllHandler(w http.ResponseWriter, req *http.Request) {
	status, _ := fleetClient.MachineAll()
	renderer.JSON(w, http.StatusOK, status)
}

func statusAllHandler(w http.ResponseWriter, req *http.Request) {
	status, _ := fleetClient.StatusAll()
	renderer.JSON(w, http.StatusOK, status)
}

func statusHandler(w http.ResponseWriter, req *http.Request) {
	key := mux.Vars(req)["id"]
	status, _ := fleetClient.StatusUnit(key)
	renderer.JSON(w, http.StatusOK, status)
}

// websocket handler

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// write journal message
	key := mux.Vars(r)["id"]
	output, _ := fleetClient.JournalF(key)
	for line := range output {
		conn.WriteMessage(websocket.TextMessage, []byte(line))
	}
}
<|endoftext|>


<commit_before>package scroll

import (
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gorilla/mux"
	"github.com/mailgun/log"
	"github.com/mailgun/manners"
	"github.com/mailgun/metrics"

	"github.com/mailgun/scroll/registry"
	"github.com/mailgun/scroll/vulcan/middleware"
)

const (
	// Suggested result set limit for APIs that may return many entries (e.g. paging).
	DefaultLimit = 100

	// Suggested max allowed result set limit for APIs that may return many entries (e.g. paging).
	MaxLimit = 10000

	// Suggested max allowed amount of entries that batch APIs can accept (e.g. batch uploads).
	MaxBatchSize = 1000

	// Interval between Vulcand heartbeats (if the app if configured to register in it).
	defaultRegisterInterval = 2 * time.Second
)

// Represents an app.
type App struct {
	Config      AppConfig
	router      *mux.Router
	stats       *appStats
	heartbeater *registry.Heartbeater
}

// Represents a configuration object an app is created with.
type AppConfig struct {
	// name of the app being created
	Name string

	// IP/port the app will bind to
	ListenIP   string
	ListenPort int

	// optional router to use
	Router *mux.Router

	// hostnames of the public and protected API entrypoints used for vulcan registration
	PublicAPIHost    string
	ProtectedAPIHost string
	ProtectedAPIURL  string

	// how to register the app's endpoint and handlers in vulcan
	Registry registry.Registry
	Interval time.Duration

	// metrics service used for emitting the app's real-time metrics
	Client metrics.Client
}

// Create a new app.
func NewApp() *App {
	return NewAppWithConfig(AppConfig{})
}

// Create a new app with the provided configuration.
func NewAppWithConfig(config AppConfig) *App {
	router := config.Router
	if router == nil {
		router = mux.NewRouter()
	}

	interval := config.Interval
	if interval == 0 {
		interval = defaultRegisterInterval
	}

	registration := &registry.AppRegistration{Name: config.Name, Host: config.ListenIP, Port: config.ListenPort}
	heartbeater := registry.NewHeartbeater(registration, config.Registry, interval)

	return &App{
		Config:      config,
		router:      router,
		heartbeater: heartbeater,
		stats:       newAppStats(config.Client),
	}
}

// Register a handler function.
//
// If vulcan registration is enabled in the both app config and handler spec,
// the handler will be registered in the local etcd instance.
func (app *App) AddHandler(spec Spec) error {
	var handler http.HandlerFunc

	// make a handler depending on the function provided in the spec
	if spec.RawHandler != nil {
		handler = spec.RawHandler
	} else if spec.Handler != nil {
		handler = MakeHandler(app, spec.Handler, spec)
	} else if spec.HandlerWithBody != nil {
		handler = MakeHandlerWithBody(app, spec.HandlerWithBody, spec)
	} else {
		return fmt.Errorf("the spec does not provide a handler function: %v", spec)
	}

	for _, path := range spec.Paths {
		// register a handler in the router
		route := app.router.HandleFunc(path, handler).Methods(spec.Methods...)
		if len(spec.Headers) != 0 {
			route.Headers(spec.Headers...)
		}

		app.registerLocation(spec.Methods, path, spec.Scopes, spec.Middlewares)
	}

	return nil
}

// GetHandler returns HTTP compatible Handler interface.
func (app *App) GetHandler() http.Handler {
	return app.router
}

// SetNotFoundHandler sets the handler for the case when URL can not be matched by the router.
func (app *App) SetNotFoundHandler(fn http.HandlerFunc) {
	app.router.NotFoundHandler = fn
}

// IsPublicRequest determines whether the provided request came through the public HTTP endpoint.
func (app *App) IsPublicRequest(request *http.Request) bool {
	return request.Host == app.Config.PublicAPIHost
}

// Start the app on the configured host/port.
//
// Supports graceful shutdown on 'kill' and 'int' signals.
func (app *App) Run() error {
	http.Handle("/", app.router)

	// toggle heartbeat on SIGUSR1
	go func() {
		app.heartbeater.Start()
		heartbeatChan := make(chan os.Signal, 1)
		signal.Notify(heartbeatChan, syscall.SIGUSR1)

		for s := range heartbeatChan {
			log.Infof("Received signal: %v, toggling heartbeat", s)
			app.heartbeater.Toggle()
		}
	}()

	// listen for a shutdown signal
	go func() {
		exitChan := make(chan os.Signal, 1)
		signal.Notify(exitChan, os.Interrupt, os.Kill)
		s := <-exitChan
		log.Infof("Got shutdown signal: %v", s)
		manners.Close()
	}()

	addr := fmt.Sprintf("%v:%v", app.Config.ListenIP, app.Config.ListenPort)
	return manners.ListenAndServe(addr, nil)
}

// registerLocation is a helper for registering handlers in vulcan.
func (app *App) registerLocation(methods []string, path string, scopes []Scope, middlewares []middleware.Middleware) {
	for _, scope := range scopes {
		app.registerLocationForScope(methods, path, scope, middlewares)
	}
}

// registerLocationForScope registers a location with a specified scope.
func (app *App) registerLocationForScope(methods []string, path string, scope Scope, middlewares []middleware.Middleware) {
	host, err := app.apiHostForScope(scope)
	if err != nil {
		log.Errorf("Failed to register a location: %v", err)
		return
	}
	app.registerLocationForHost(methods, path, host, middlewares)
}

// registerLocationForHost registers a location for a specified hostname.
func (app *App) registerLocationForHost(methods []string, path, host string, middlewares []middleware.Middleware) {
	r := &registry.HandlerRegistration{
		Name:        app.Config.Name,
		Host:        host,
		Path:        path,
		Methods:     methods,
		Middlewares: middlewares,
	}
	app.Config.Registry.RegisterHandler(r)

	log.Infof("Registered: %v", r)
}

// apiHostForScope is a helper that returns an appropriate API hostname for a provided scope.
func (app *App) apiHostForScope(scope Scope) (string, error) {
	if scope == ScopePublic {
		return app.Config.PublicAPIHost, nil
	} else if scope == ScopeProtected {
		return app.Config.ProtectedAPIHost, nil
	} else {
		return "", fmt.Errorf("unknown scope value: %v", scope)
	}
}
<commit_msg>Register sigterm for graceful shutdown<commit_after>package scroll

import (
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gorilla/mux"
	"github.com/mailgun/log"
	"github.com/mailgun/manners"
	"github.com/mailgun/metrics"

	"github.com/mailgun/scroll/registry"
	"github.com/mailgun/scroll/vulcan/middleware"
)

const (
	// Suggested result set limit for APIs that may return many entries (e.g. paging).
	DefaultLimit = 100

	// Suggested max allowed result set limit for APIs that may return many entries (e.g. paging).
	MaxLimit = 10000

	// Suggested max allowed amount of entries that batch APIs can accept (e.g. batch uploads).
	MaxBatchSize = 1000

	// Interval between Vulcand heartbeats (if the app if configured to register in it).
	defaultRegisterInterval = 2 * time.Second
)

// Represents an app.
type App struct {
	Config      AppConfig
	router      *mux.Router
	stats       *appStats
	heartbeater *registry.Heartbeater
}

// Represents a configuration object an app is created with.
type AppConfig struct {
	// name of the app being created
	Name string

	// IP/port the app will bind to
	ListenIP   string
	ListenPort int

	// optional router to use
	Router *mux.Router

	// hostnames of the public and protected API entrypoints used for vulcan registration
	PublicAPIHost    string
	ProtectedAPIHost string
	ProtectedAPIURL  string

	// how to register the app's endpoint and handlers in vulcan
	Registry registry.Registry
	Interval time.Duration

	// metrics service used for emitting the app's real-time metrics
	Client metrics.Client
}

// Create a new app.
func NewApp() *App {
	return NewAppWithConfig(AppConfig{})
}

// Create a new app with the provided configuration.
func NewAppWithConfig(config AppConfig) *App {
	router := config.Router
	if router == nil {
		router = mux.NewRouter()
	}

	interval := config.Interval
	if interval == 0 {
		interval = defaultRegisterInterval
	}

	registration := &registry.AppRegistration{Name: config.Name, Host: config.ListenIP, Port: config.ListenPort}
	heartbeater := registry.NewHeartbeater(registration, config.Registry, interval)

	return &App{
		Config:      config,
		router:      router,
		heartbeater: heartbeater,
		stats:       newAppStats(config.Client),
	}
}

// Register a handler function.
//
// If vulcan registration is enabled in the both app config and handler spec,
// the handler will be registered in the local etcd instance.
func (app *App) AddHandler(spec Spec) error {
	var handler http.HandlerFunc

	// make a handler depending on the function provided in the spec
	if spec.RawHandler != nil {
		handler = spec.RawHandler
	} else if spec.Handler != nil {
		handler = MakeHandler(app, spec.Handler, spec)
	} else if spec.HandlerWithBody != nil {
		handler = MakeHandlerWithBody(app, spec.HandlerWithBody, spec)
	} else {
		return fmt.Errorf("the spec does not provide a handler function: %v", spec)
	}

	for _, path := range spec.Paths {
		// register a handler in the router
		route := app.router.HandleFunc(path, handler).Methods(spec.Methods...)
		if len(spec.Headers) != 0 {
			route.Headers(spec.Headers...)
		}

		app.registerLocation(spec.Methods, path, spec.Scopes, spec.Middlewares)
	}

	return nil
}

// GetHandler returns HTTP compatible Handler interface.
func (app *App) GetHandler() http.Handler {
	return app.router
}

// SetNotFoundHandler sets the handler for the case when URL can not be matched by the router.
func (app *App) SetNotFoundHandler(fn http.HandlerFunc) {
	app.router.NotFoundHandler = fn
}

// IsPublicRequest determines whether the provided request came through the public HTTP endpoint.
func (app *App) IsPublicRequest(request *http.Request) bool {
	return request.Host == app.Config.PublicAPIHost
}

// Start the app on the configured host/port.
//
// Supports graceful shutdown on 'kill' and 'int' signals.
func (app *App) Run() error {
	http.Handle("/", app.router)

	// toggle heartbeat on SIGUSR1
	go func() {
		app.heartbeater.Start()
		heartbeatChan := make(chan os.Signal, 1)
		signal.Notify(heartbeatChan, syscall.SIGUSR1)

		for s := range heartbeatChan {
			log.Infof("Received signal: %v, toggling heartbeat", s)
			app.heartbeater.Toggle()
		}
	}()

	// listen for a shutdown signal
	go func() {
		exitChan := make(chan os.Signal, 1)
		signal.Notify(exitChan, syscall.SIGINT, syscall.SIGTERM)
		s := <-exitChan
		log.Infof("Got shutdown signal: %v", s)
		manners.Close()
	}()

	addr := fmt.Sprintf("%v:%v", app.Config.ListenIP, app.Config.ListenPort)
	return manners.ListenAndServe(addr, nil)
}

// registerLocation is a helper for registering handlers in vulcan.
func (app *App) registerLocation(methods []string, path string, scopes []Scope, middlewares []middleware.Middleware) {
	for _, scope := range scopes {
		app.registerLocationForScope(methods, path, scope, middlewares)
	}
}

// registerLocationForScope registers a location with a specified scope.
func (app *App) registerLocationForScope(methods []string, path string, scope Scope, middlewares []middleware.Middleware) {
	host, err := app.apiHostForScope(scope)
	if err != nil {
		log.Errorf("Failed to register a location: %v", err)
		return
	}
	app.registerLocationForHost(methods, path, host, middlewares)
}

// registerLocationForHost registers a location for a specified hostname.
func (app *App) registerLocationForHost(methods []string, path, host string, middlewares []middleware.Middleware) {
	r := &registry.HandlerRegistration{
		Name:        app.Config.Name,
		Host:        host,
		Path:        path,
		Methods:     methods,
		Middlewares: middlewares,
	}
	app.Config.Registry.RegisterHandler(r)

	log.Infof("Registered: %v", r)
}

// apiHostForScope is a helper that returns an appropriate API hostname for a provided scope.
func (app *App) apiHostForScope(scope Scope) (string, error) {
	if scope == ScopePublic {
		return app.Config.PublicAPIHost, nil
	} else if scope == ScopeProtected {
		return app.Config.ProtectedAPIHost, nil
	} else {
		return "", fmt.Errorf("unknown scope value: %v", scope)
	}
}
<|endoftext|>


<commit_before>package mdqi

import (
	"fmt"
	"io"
	"os"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/peterh/liner"
	"github.com/pkg/errors"
)

var (
	ErrSlashCommandNotFound = errors.New("unknown SlashCommand")
	ErrNotASlashCommand     = errors.New("there are no SlashCommand")
)

type App struct {
	// Alive turns into false, mdqi will exit.
	Alive bool

	// cmdPath is path to mdq command.
	cmdPath string

	// historyPath is path to command history file for liner.
	historyPath string

	// slashCommandDefinition holds SlashCommandDefinition.
	// app.slashCommandDefinition[category][name] = SlashCommandDefinition
	slashCommandDefinition map[string]map[string]SlashCommandDefinition
}

type Conf struct {
}

type Result struct {
	Database string
	Columns  []string
	Rows     []map[string]interface{}
}

type SlashCommand struct {
	Category string
	Name     string
	Args     []string
}

type SlashCommandDefinition struct {
	Category string
	Name     string
	Handler  SlashCommandHandler
}

func NewApp(conf Conf) (*App, error) {
	// TODO: Check if mdq command exists by exec.LookPath.
	// TODO: Make historyPath configuarable.

	historyPath, err := createHistoryFile()
	if err != nil {
		return nil, errors.Wrap(err, "failed to create history file")
	}

	app := &App{
		Alive: true,

		cmdPath:                "mdq",
		historyPath:            historyPath,
		slashCommandDefinition: map[string]map[string]SlashCommandDefinition{},
	}

	app.initSlashCommands()

	return app, nil
}

func createHistoryFile() (string, error) {
	usr, err := user.Current()
	if err != nil {
		return "", errors.Wrap(err, "failed to get current user")
	}

	path := filepath.Join(usr.HomeDir, ".mdqi_history")

	if _, err := os.Stat(path); os.IsNotExist(err) {
		if _, err := os.Create(path); err != nil {
			return "", errors.Wrap(err, "failed to create history file")
		}
	}

	return path, nil
}

func (app *App) initSlashCommands() {
	app.RegisterSlashCommand("exit", "", SlashCommandExit)
}

func (app *App) Run() {
	app.runLiner()
}

func (app *App) runLiner() {
	line := liner.NewLiner()
	defer line.Close()

	line.SetCtrlCAborts(true)

	if f, err := os.Open(app.historyPath); err == nil {
		line.ReadHistory(f)
		f.Close()
	} else {
		fmt.Fprintln(os.Stderr, "failed to read command history: ", err)
	}

LOOP:
	for {
		if !app.Alive {
			fmt.Println("bye")
			break LOOP
		}

		l, err := line.Prompt("mdq> ")

		switch err {
		case nil:
			scmd, _ := ParseSlashCommand(l)
			if scmd != nil {
				app.runSlashCommand(scmd)
				continue
			}

			results, err := app.RunCmd(strings.Trim(l, " \n"))
			if err != nil {
				fmt.Fprintln(os.Stderr, err.Error())
			}

			Print(results)

			line.AppendHistory(l)
		case liner.ErrPromptAborted:
			fmt.Fprintln(os.Stderr, "aborted")
			break LOOP
		case io.EOF:
			fmt.Println("bye")
			break LOOP
		default:
			fmt.Fprintln(os.Stderr, "error on reading line: ", err)
			break LOOP
		}

		if f, err := os.Create(app.historyPath); err == nil {
			if _, err := line.WriteHistory(f); err != nil {
				fmt.Fprintln(os.Stderr, "failed to write history: ", err)
			}

			f.Close()
		} else {
			fmt.Fprintln(os.Stderr, "failed to create history file: ", err)
		}
	}
}

func (app *App) runSlashCommand(scmd *SlashCommand) {
	sdef, err := app.FindSlashCommandDefinition(scmd.Category, scmd.Name)

	switch err {
	case nil:
		if err := sdef.Handler(app, scmd); err != nil {
			fmt.Fprintln(os.Stderr, "failed to handle slash command:", err)
		}
	case ErrSlashCommandNotFound:
		fmt.Fprintln(os.Stderr, "unknown slash command")
	}

	return
}
<commit_msg>refactor: init and save history as functions.<commit_after>package mdqi

import (
	"fmt"
	"io"
	"os"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/peterh/liner"
	"github.com/pkg/errors"
)

var (
	ErrSlashCommandNotFound = errors.New("unknown SlashCommand")
	ErrNotASlashCommand     = errors.New("there are no SlashCommand")
)

type App struct {
	// Alive turns into false, mdqi will exit.
	Alive bool

	// cmdPath is path to mdq command.
	cmdPath string

	// historyPath is path to command history file for liner.
	historyPath string

	// slashCommandDefinition holds SlashCommandDefinition.
	// app.slashCommandDefinition[category][name] = SlashCommandDefinition
	slashCommandDefinition map[string]map[string]SlashCommandDefinition
}

type Conf struct {
}

type Result struct {
	Database string
	Columns  []string
	Rows     []map[string]interface{}
}

type SlashCommand struct {
	Category string
	Name     string
	Args     []string
}

type SlashCommandDefinition struct {
	Category string
	Name     string
	Handler  SlashCommandHandler
}

func NewApp(conf Conf) (*App, error) {
	// TODO: Check if mdq command exists by exec.LookPath.
	// TODO: Make historyPath configuarable.

	historyPath, err := createHistoryFile()
	if err != nil {
		return nil, errors.Wrap(err, "failed to create history file")
	}

	app := &App{
		Alive: true,

		cmdPath:                "mdq",
		historyPath:            historyPath,
		slashCommandDefinition: map[string]map[string]SlashCommandDefinition{},
	}

	app.initSlashCommands()

	return app, nil
}

func createHistoryFile() (string, error) {
	usr, err := user.Current()
	if err != nil {
		return "", errors.Wrap(err, "failed to get current user")
	}

	path := filepath.Join(usr.HomeDir, ".mdqi_history")

	if _, err := os.Stat(path); os.IsNotExist(err) {
		if _, err := os.Create(path); err != nil {
			return "", errors.Wrap(err, "failed to create history file")
		}
	}

	return path, nil
}

func (app *App) initSlashCommands() {
	app.RegisterSlashCommand("exit", "", SlashCommandExit)
}

func (app *App) Run() {
	app.runLiner()
}

func (app *App) runLiner() {
	line := liner.NewLiner()
	defer line.Close()

	line.SetCtrlCAborts(true)

	app.initHistory(line)

LOOP:
	for {
		if !app.Alive {
			fmt.Println("bye")
			break LOOP
		}

		l, err := line.Prompt("mdq> ")

		switch err {
		case nil:
			scmd, _ := ParseSlashCommand(l)
			if scmd != nil {
				app.runSlashCommand(scmd)
				continue
			}

			results, err := app.RunCmd(strings.Trim(l, " \n"))
			if err != nil {
				fmt.Fprintln(os.Stderr, err.Error())
			}

			Print(results)

			line.AppendHistory(l)
		case liner.ErrPromptAborted:
			fmt.Fprintln(os.Stderr, "aborted")
			break LOOP
		case io.EOF:
			fmt.Println("bye")
			break LOOP
		default:
			fmt.Fprintln(os.Stderr, "error on reading line: ", err)
			break LOOP
		}

		app.saveHistory(line)
	}
}

func (app *App) initHistory(line *liner.State) {
	if f, err := os.Open(app.historyPath); err == nil {
		line.ReadHistory(f)
		f.Close()
	} else {
		fmt.Fprintln(os.Stderr, "failed to read command history: ", err)
	}
}

func (app *App) saveHistory(line *liner.State) {
	if f, err := os.Create(app.historyPath); err == nil {
		if _, err := line.WriteHistory(f); err != nil {
			fmt.Fprintln(os.Stderr, "failed to write history: ", err)
		}

		f.Close()
	} else {
		fmt.Fprintln(os.Stderr, "failed to create history file: ", err)
	}
}

func (app *App) runSlashCommand(scmd *SlashCommand) {
	sdef, err := app.FindSlashCommandDefinition(scmd.Category, scmd.Name)

	switch err {
	case nil:
		if err := sdef.Handler(app, scmd); err != nil {
			fmt.Fprintln(os.Stderr, "failed to handle slash command:", err)
		}
	case ErrSlashCommandNotFound:
		fmt.Fprintln(os.Stderr, "unknown slash command")
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"

	"github.com/jmoiron/sqlx"
	"github.com/training_project/config"
	"github.com/training_project/controller/review"
	"github.com/training_project/database"
	reviewModel "github.com/training_project/model/review"

	logging "gopkg.in/tokopedia/logging.v1"
)

var cfg config.Config

func init() {
	// get config from database.ini
	// assigne to global variable cfg
	ok := logging.ReadModuleConfig(&cfg, "/etc/test", "test") || logging.ReadModuleConfig(&cfg, "config", "test")
	if !ok {
		log.Fatalln("failed to read config")
	}
}

func main() {
	//getting list of all the connection.
	listConnection := database.SystemConnection()

	//getting redis connection convert it from interface to *redisClient.
	//redisConn := listConnection["redis"].(*redis.Client)

	// get postgre connection.
	postgreConn := listConnection["postgre"].(*sqlx.DB)

	//pass to model
	reviewData := &reviewModel.ReviewData{}
	reviewData.GetConn(postgreConn)

	http.HandleFunc("/", review.CheckDataExist)

	port := ":8080"
	fmt.Println("App Started on port = ", port)
	err := http.ListenAndServe(port, nil)
	if err != nil {
		log.Panic("App Started Failed = ", err.Error())
	}

}
<commit_msg>create endpoint for driver<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"

	"github.com/jmoiron/sqlx"
	"github.com/training_project/config"
	"github.com/training_project/controller/review"
	"github.com/training_project/database"
	reviewModel "github.com/training_project/model/review"
	"github.com/training_project/util/logger"

	logging "gopkg.in/tokopedia/logging.v1"
)

var cfg config.Config

func init() {
	// get config from database.ini
	// assigne to global variable cfg
	ok := logging.ReadModuleConfig(&cfg, "/etc/test", "test") || logging.ReadModuleConfig(&cfg, "config", "test")
	if !ok {
		log.Fatalln("failed to read config")
	}

	logger.InitLogger("App :: ", "./logs", "App.txt")
}

func main() {
	//getting list of all the connection.
	listConnection := database.SystemConnection()

	//getting redis connection convert it from interface to *redisClient.
	//redisConn := listConnection["redis"].(*redis.Client)

	// get postgre connection.
	postgreConn := listConnection["postgre"].(*sqlx.DB)

	//pass to model
	reviewData := &reviewModel.ReviewData{}
	reviewData.GetConn(postgreConn)

	http.HandleFunc("/", review.CheckDataExist)
	http.HandleFunc("/driver", review.CheckDataExist)

	port := ":8080"
	fmt.Println("App Started on port = ", port)
	err := http.ListenAndServe(port, nil)
	if err != nil {
		log.Panic("App Started Failed = ", err.Error())
	}

}
<|endoftext|>


<commit_before>package hacknews

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
)

type Initializer struct {
	Story   string
	NbPosts int
}

type Post struct {
	Id          int    `json:"id"`
	Deleted     bool   `json:"deleted"`
	Type        string `json:"type"`
	By          string `json:"by"`
	Time        int    `json:"time"`
	Text        string `json:"text"`
	Dead        bool   `json:"dead"`
	Parent      int    `json:"parent"`
	Poll        int    `json:"poll"`
	Kids        []int  `json:"kids"`
	Url         string `json:"url"`
	Score       int    `json:"score"`
	Title       string `json:"title"`
	Parts       []int  `json:"parts"`
	Descendants int    `json:"descendants"`
}

// Return the ids of a story
func (init Initializer) GetCodesStory() ([]int, error) {
	resp, errFetch := http.Get("https://hacker-news.firebaseio.com/v0/" + init.Story + ".json?print=pretty")
	if errFetch != nil {
		return nil, fmt.Errorf("Error while fetching codes story : %v", errFetch)
	}
	defer resp.Body.Close()
	body, errRead := ioutil.ReadAll(resp.Body)
	if errRead != nil {
		return nil, fmt.Errorf("Error while reading Body of codes story : %v", errRead)
	}
	jsonDecoded := []int{}
	errDecoded := json.Unmarshal(body, &jsonDecoded)
	if errDecoded != nil {
		return nil, fmt.Errorf("Error while decoding json from codes story : %v", errDecoded)
	}
	return jsonDecoded, nil
}

// Return the posts of story thanks their ids
func (init Initializer) GetPostStory(codes []int) ([]Post, error) {
	post := Post{}
	myData := []Post{}
	for i := 0; i < init.NbPosts; i++ {
		resp, errFetch := http.Get("https://hacker-news.firebaseio.com/v0/item/" + strconv.Itoa(codes[i]) + ".json?print=pretty")
		if errFetch != nil {
			return nil, fmt.Errorf("Error while fetching posts story : %v", errFetch)
		}
		body, errRead := ioutil.ReadAll(resp.Body)
		if errRead != nil {
			return nil, fmt.Errorf("Error while reading Body of posts story : %v", errRead)
		}
		errDecoded := json.Unmarshal(body, &post)
		if errDecoded != nil {
			return nil, fmt.Errorf("Error while decoding json from posts story : %v", errDecoded)
		}
		myData = append(myData, post)
	}
	return myData, nil
}
<commit_msg>Fix godoc<commit_after>package hacknews

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
)

// Initializer represents the kind of requests that we want to fetch
type Initializer struct {
	Story   string
	NbPosts int
}

// Post represents the json object returned by the API
type Post struct {
	Id          int    `json:"id"`
	Deleted     bool   `json:"deleted"`
	Type        string `json:"type"`
	By          string `json:"by"`
	Time        int    `json:"time"`
	Text        string `json:"text"`
	Dead        bool   `json:"dead"`
	Parent      int    `json:"parent"`
	Poll        int    `json:"poll"`
	Kids        []int  `json:"kids"`
	Url         string `json:"url"`
	Score       int    `json:"score"`
	Title       string `json:"title"`
	Parts       []int  `json:"parts"`
	Descendants int    `json:"descendants"`
}

// Return the ids of a story
func (init Initializer) GetCodesStory() ([]int, error) {
	resp, errFetch := http.Get("https://hacker-news.firebaseio.com/v0/" + init.Story + ".json?print=pretty")
	if errFetch != nil {
		return nil, fmt.Errorf("Error while fetching codes story : %v", errFetch)
	}
	defer resp.Body.Close()
	body, errRead := ioutil.ReadAll(resp.Body)
	if errRead != nil {
		return nil, fmt.Errorf("Error while reading Body of codes story : %v", errRead)
	}
	jsonDecoded := []int{}
	errDecoded := json.Unmarshal(body, &jsonDecoded)
	if errDecoded != nil {
		return nil, fmt.Errorf("Error while decoding json from codes story : %v", errDecoded)
	}
	return jsonDecoded, nil
}

// Return the posts of story thanks their ids
func (init Initializer) GetPostStory(codes []int) ([]Post, error) {
	post := Post{}
	myData := []Post{}
	for i := 0; i < init.NbPosts; i++ {
		resp, errFetch := http.Get("https://hacker-news.firebaseio.com/v0/item/" + strconv.Itoa(codes[i]) + ".json?print=pretty")
		if errFetch != nil {
			return nil, fmt.Errorf("Error while fetching posts story : %v", errFetch)
		}
		body, errRead := ioutil.ReadAll(resp.Body)
		if errRead != nil {
			return nil, fmt.Errorf("Error while reading Body of posts story : %v", errRead)
		}
		errDecoded := json.Unmarshal(body, &post)
		if errDecoded != nil {
			return nil, fmt.Errorf("Error while decoding json from posts story : %v", errDecoded)
		}
		myData = append(myData, post)
	}
	return myData, nil
}
<|endoftext|>


<commit_before>package nsqlookupd

import (
	"bufio"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/absolute8511/nsq/internal/protocol"
	"github.com/absolute8511/nsq/internal/version"
)

const (
	OLD_VERSION_PID = -11
)

type LookupProtocolV1 struct {
	ctx *Context
}

func (p *LookupProtocolV1) IOLoop(conn net.Conn) error {
	var err error
	var line string
	var zeroTime time.Time
	to := p.ctx.nsqlookupd.opts.NsqdPingTimeout

	client := NewClientV1(conn)
	reader := bufio.NewReader(client)
	for {
		if to > 0 {
			client.SetReadDeadline(time.Now().Add(to))
		} else {
			client.SetReadDeadline(zeroTime)
		}
		line, err = reader.ReadString('\n')
		if err != nil {
			break
		}

		line = strings.TrimSpace(line)
		params := strings.Split(line, " ")

		var response []byte
		response, err = p.Exec(client, reader, params)
		if err != nil {
			ctx := ""
			if parentErr := err.(protocol.ChildErr).Parent(); parentErr != nil {
				ctx = " - " + parentErr.Error()
			}
			nsqlookupLog.LogErrorf(" [%s] - %s%s", client, err, ctx)

			_, sendErr := protocol.SendResponse(client, []byte(err.Error()))
			if sendErr != nil {
				nsqlookupLog.LogErrorf(" [%s] - %s%s", client, sendErr, ctx)
				break
			}

			// errors of type FatalClientErr should forceably close the connection
			if _, ok := err.(*protocol.FatalClientErr); ok {
				break
			}
			continue
		}

		if response != nil {
			_, err = protocol.SendResponse(client, response)
			if err != nil {
				break
			}
		}
	}

	nsqlookupLog.Logf("CLIENT(%s): closing, %v", client, err)
	if client.peerInfo != nil {
		p.ctx.nsqlookupd.DB.RemoveAllByPeerId(client.peerInfo.Id)
	}
	conn.Close()
	return err
}

func (p *LookupProtocolV1) Exec(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	switch params[0] {
	case "PING":
		return p.PING(client, params)
	case "IDENTIFY":
		return p.IDENTIFY(client, reader, params[1:])
	case "REGISTER":
		return p.REGISTER(client, reader, params[1:])
	case "UNREGISTER":
		return p.UNREGISTER(client, reader, params[1:])
	}
	return nil, protocol.NewFatalClientErr(nil, "E_INVALID", fmt.Sprintf("invalid command %s", params[0]))
}

func getTopicChanFromOld(command string, params []string) (string, string, error) {
	if len(params) <= 0 {
		return "", "", protocol.NewFatalClientErr(nil, "E_INVALID", fmt.Sprintf("%s insufficient number of params", command))
	}

	topicName := params[0]
	var channelName string
	if len(params) >= 2 {
		channelName = params[1]
	}

	if !protocol.IsValidTopicName(topicName) {
		return "", "", protocol.NewFatalClientErr(nil, "E_BAD_TOPIC", fmt.Sprintf("%s topic name '%s' is not valid", command, topicName))
	}

	if channelName != "" && !protocol.IsValidChannelName(channelName) {
		return "", "", protocol.NewFatalClientErr(nil, "E_BAD_CHANNEL", fmt.Sprintf("%s channel name '%s' is not valid", command, channelName))
	}
	return topicName, channelName, nil
}

func getTopicChan(command string, params []string) (string, string, string, error) {
	if len(params) <= 1 {
		return "", "", "", protocol.NewFatalClientErr(nil, "E_INVALID", fmt.Sprintf("%s insufficient number of params", command))
	}

	topicName := params[0]
	partitionID := params[1]
	var channelName string
	if len(params) >= 3 {
		channelName = params[2]
	}

	if !protocol.IsValidTopicName(topicName) {
		return "", "", "", protocol.NewFatalClientErr(nil, "E_BAD_TOPIC", fmt.Sprintf("%s topic name '%s' is not valid", command, topicName))
	}

	if channelName != "" && !protocol.IsValidChannelName(channelName) {
		return "", "", "", protocol.NewFatalClientErr(nil, "E_BAD_CHANNEL", fmt.Sprintf("%s channel name '%s' is not valid", command, channelName))
	}

	if _, err := GetValidPartitionID(partitionID); err != nil {
		return "", "", "", protocol.NewFatalClientErr(nil, "E_BAD_PARTITIONID", fmt.Sprintf("%s partition id '%s' is not valid", command, partitionID))
	}
	return topicName, channelName, partitionID, nil
}

func (p *LookupProtocolV1) REGISTER(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	if client.peerInfo == nil {
		return nil, protocol.NewFatalClientErr(nil, "E_INVALID", "client must IDENTIFY")
	}

	topic, channel, pid, err := getTopicChan("REGISTER", params)
	if err != nil {
		// check if old nsqd
		if client.peerInfo.IsOldPeer() {
			nsqlookupLog.Logf("client %v is old node trying register", client)
			topic, channel, err = getTopicChanFromOld("REGISTER", params)
			pid = strconv.Itoa(OLD_VERSION_PID)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}

	atomic.StoreInt64(&client.peerInfo.lastUpdate, time.Now().UnixNano())
	if channel != "" {
		key := ChannelReg{
			PartitionID: pid,
			PeerId:      client.peerInfo.Id,
			Channel:     channel,
		}
		if p.ctx.nsqlookupd.DB.AddChannelReg(topic, key) {
			nsqlookupLog.Logf("DB: client(%s) REGISTER new channel: topic:%s channel:%s pid:%s",
				client, topic, channel, pid)
		}
	}
	if p.ctx.nsqlookupd.DB.AddTopicProducer(topic, pid, &Producer{peerInfo: client.peerInfo}) {
		nsqlookupLog.Logf("DB: client(%s) REGISTER new topic:%s pid:%s",
			client, topic, pid)
	}

	return []byte("OK"), nil
}

func (p *LookupProtocolV1) UNREGISTER(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	if client.peerInfo == nil {
		return nil, protocol.NewFatalClientErr(nil, "E_INVALID", "client must IDENTIFY")
	}

	topic, channel, pid, err := getTopicChan("UNREGISTER", params)
	if err != nil {
		// check if old nsqd
		if client.peerInfo.IsOldPeer() {
			nsqlookupLog.Logf("client %v is old node trying unregister", client)
			pid = strconv.Itoa(OLD_VERSION_PID)
			topic, channel, err = getTopicChanFromOld("UNREGISTER", params)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}

	atomic.StoreInt64(&client.peerInfo.lastUpdate, time.Now().UnixNano())
	if channel != "" {
		key := ChannelReg{
			PartitionID: pid,
			PeerId:      client.peerInfo.Id,
			Channel:     channel,
		}

		removed := p.ctx.nsqlookupd.DB.RemoveChannelReg(topic, key)
		if removed {
			nsqlookupLog.Logf("DB: client(%s) UNREGISTER channel %v on topic:%s-%v",
				client, channel, topic, pid)
		}
	} else {
		// no channel was specified so this is a topic unregistration
		// remove all of the channel registrations...
		// normally this shouldn't happen which is why we print a warning message
		// if anything is actually removed
		key := ChannelReg{
			PartitionID: pid,
			PeerId:      client.peerInfo.Id,
			Channel:     "*",
		}

		removed := p.ctx.nsqlookupd.DB.RemoveChannelReg(topic, key)
		if removed {
			nsqlookupLog.LogWarningf(" client(%s) unexpected UNREGISTER all channels under topic:%s pid:%s",
				client, topic, pid)
		}

		if removed := p.ctx.nsqlookupd.DB.RemoveTopicProducer(topic, pid, client.peerInfo.Id); removed {
			nsqlookupLog.Logf("DB: client(%s) UNREGISTER topic :%s pid:%s",
				client, topic, pid)
		}
	}

	return []byte("OK"), nil
}

func (p *LookupProtocolV1) IDENTIFY(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	var err error

	if client.peerInfo != nil {
		return nil, protocol.NewFatalClientErr(err, "E_INVALID", "cannot IDENTIFY again")
	}

	var bodyLen int32
	err = binary.Read(reader, binary.BigEndian, &bodyLen)
	if err != nil {
		return nil, protocol.NewFatalClientErr(err, "E_BAD_BODY", "IDENTIFY failed to read body size")
	}

	body := make([]byte, bodyLen)
	_, err = io.ReadFull(reader, body)
	if err != nil {
		return nil, protocol.NewFatalClientErr(err, "E_BAD_BODY", "IDENTIFY failed to read body")
	}

	// body is a json structure with producer information
	// Id should be identified by the client.
	peerInfo := PeerInfo{Id: ""}
	err = json.Unmarshal(body, &peerInfo)
	if err != nil {
		return nil, protocol.NewFatalClientErr(err, "E_BAD_BODY", "IDENTIFY failed to decode JSON body")
	}

	peerInfo.RemoteAddress = client.RemoteAddr().String()
	peerInfo.Id = peerInfo.RemoteAddress

	// require all fields
	if peerInfo.Id == "" || peerInfo.BroadcastAddress == "" || peerInfo.TCPPort == 0 || peerInfo.HTTPPort == 0 || peerInfo.Version == "" {
		nsqlookupLog.Logf("identify info missing: %v", peerInfo)
		return nil, protocol.NewFatalClientErr(nil, "E_BAD_BODY", "IDENTIFY missing fields")
	}

	if p.ctx.nsqlookupd.RealTCPAddr() == nil || p.ctx.nsqlookupd.RealHTTPAddr() == nil {
		nsqlookupLog.Logf("client(%s) register before the server is ready", client)
		return nil, protocol.NewFatalClientErr(nil, "E_NOT_READY", "The server is not ready for use")
	}

	atomic.StoreInt64(&peerInfo.lastUpdate, time.Now().UnixNano())

	nsqlookupLog.Logf("CLIENT(%s): IDENTIFY peer: %v ",
		client, peerInfo)

	client.peerInfo = &peerInfo
	if p.ctx.nsqlookupd.DB.addPeerClient(peerInfo.Id, client.peerInfo) {
		nsqlookupLog.Logf("DB: client(%s) REGISTER new peer", client)
	}

	// build a response
	data := make(map[string]interface{})
	data["tcp_port"] = p.ctx.nsqlookupd.RealTCPAddr().Port
	data["http_port"] = p.ctx.nsqlookupd.RealHTTPAddr().Port
	data["version"] = version.Binary
	hostname, err := os.Hostname()
	if err != nil {
		nsqlookupLog.LogErrorf("ERROR: unable to get hostname %s", err)
	}
	data["broadcast_address"] = p.ctx.nsqlookupd.opts.BroadcastAddress
	data["hostname"] = hostname

	response, err := json.Marshal(data)
	if err != nil {
		nsqlookupLog.LogErrorf(" marshaling %v", data)
		return []byte("OK"), nil
	}
	return response, nil
}

func (p *LookupProtocolV1) PING(client *ClientV1, params []string) ([]byte, error) {
	if client.peerInfo != nil {
		// we could get a PING before other commands on the same client connection
		cur := time.Unix(0, atomic.LoadInt64(&client.peerInfo.lastUpdate))
		now := time.Now()
		nsqlookupLog.LogDebugf("CLIENT(%s): pinged (last ping %s)", client.peerInfo.Id,
			now.Sub(cur))
		atomic.StoreInt64(&client.peerInfo.lastUpdate, now.UnixNano())
	}
	return []byte("OK"), nil
}
<commit_msg>allow identify without http server<commit_after>package nsqlookupd

import (
	"bufio"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/absolute8511/nsq/internal/protocol"
	"github.com/absolute8511/nsq/internal/version"
)

const (
	OLD_VERSION_PID = -11
)

type LookupProtocolV1 struct {
	ctx *Context
}

func (p *LookupProtocolV1) IOLoop(conn net.Conn) error {
	var err error
	var line string
	var zeroTime time.Time
	to := p.ctx.nsqlookupd.opts.NsqdPingTimeout

	client := NewClientV1(conn)
	reader := bufio.NewReader(client)
	for {
		if to > 0 {
			client.SetReadDeadline(time.Now().Add(to))
		} else {
			client.SetReadDeadline(zeroTime)
		}
		line, err = reader.ReadString('\n')
		if err != nil {
			break
		}

		line = strings.TrimSpace(line)
		params := strings.Split(line, " ")

		var response []byte
		response, err = p.Exec(client, reader, params)
		if err != nil {
			ctx := ""
			if parentErr := err.(protocol.ChildErr).Parent(); parentErr != nil {
				ctx = " - " + parentErr.Error()
			}
			nsqlookupLog.LogErrorf(" [%s] - %s%s", client, err, ctx)

			_, sendErr := protocol.SendResponse(client, []byte(err.Error()))
			if sendErr != nil {
				nsqlookupLog.LogErrorf(" [%s] - %s%s", client, sendErr, ctx)
				break
			}

			// errors of type FatalClientErr should forceably close the connection
			if _, ok := err.(*protocol.FatalClientErr); ok {
				break
			}
			continue
		}

		if response != nil {
			_, err = protocol.SendResponse(client, response)
			if err != nil {
				break
			}
		}
	}

	nsqlookupLog.Logf("CLIENT(%s): closing, %v", client, err)
	if client.peerInfo != nil {
		p.ctx.nsqlookupd.DB.RemoveAllByPeerId(client.peerInfo.Id)
	}
	conn.Close()
	return err
}

func (p *LookupProtocolV1) Exec(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	switch params[0] {
	case "PING":
		return p.PING(client, params)
	case "IDENTIFY":
		return p.IDENTIFY(client, reader, params[1:])
	case "REGISTER":
		return p.REGISTER(client, reader, params[1:])
	case "UNREGISTER":
		return p.UNREGISTER(client, reader, params[1:])
	}
	return nil, protocol.NewFatalClientErr(nil, "E_INVALID", fmt.Sprintf("invalid command %s", params[0]))
}

func getTopicChanFromOld(command string, params []string) (string, string, error) {
	if len(params) <= 0 {
		return "", "", protocol.NewFatalClientErr(nil, "E_INVALID", fmt.Sprintf("%s insufficient number of params", command))
	}

	topicName := params[0]
	var channelName string
	if len(params) >= 2 {
		channelName = params[1]
	}

	if !protocol.IsValidTopicName(topicName) {
		return "", "", protocol.NewFatalClientErr(nil, "E_BAD_TOPIC", fmt.Sprintf("%s topic name '%s' is not valid", command, topicName))
	}

	if channelName != "" && !protocol.IsValidChannelName(channelName) {
		return "", "", protocol.NewFatalClientErr(nil, "E_BAD_CHANNEL", fmt.Sprintf("%s channel name '%s' is not valid", command, channelName))
	}
	return topicName, channelName, nil
}

func getTopicChan(command string, params []string) (string, string, string, error) {
	if len(params) <= 1 {
		return "", "", "", protocol.NewFatalClientErr(nil, "E_INVALID", fmt.Sprintf("%s insufficient number of params", command))
	}

	topicName := params[0]
	partitionID := params[1]
	var channelName string
	if len(params) >= 3 {
		channelName = params[2]
	}

	if !protocol.IsValidTopicName(topicName) {
		return "", "", "", protocol.NewFatalClientErr(nil, "E_BAD_TOPIC", fmt.Sprintf("%s topic name '%s' is not valid", command, topicName))
	}

	if channelName != "" && !protocol.IsValidChannelName(channelName) {
		return "", "", "", protocol.NewFatalClientErr(nil, "E_BAD_CHANNEL", fmt.Sprintf("%s channel name '%s' is not valid", command, channelName))
	}

	if _, err := GetValidPartitionID(partitionID); err != nil {
		return "", "", "", protocol.NewFatalClientErr(nil, "E_BAD_PARTITIONID", fmt.Sprintf("%s partition id '%s' is not valid", command, partitionID))
	}
	return topicName, channelName, partitionID, nil
}

func (p *LookupProtocolV1) REGISTER(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	if client.peerInfo == nil {
		return nil, protocol.NewFatalClientErr(nil, "E_INVALID", "client must IDENTIFY")
	}

	topic, channel, pid, err := getTopicChan("REGISTER", params)
	if err != nil {
		// check if old nsqd
		if client.peerInfo.IsOldPeer() {
			nsqlookupLog.Logf("client %v is old node trying register", client)
			topic, channel, err = getTopicChanFromOld("REGISTER", params)
			pid = strconv.Itoa(OLD_VERSION_PID)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}

	atomic.StoreInt64(&client.peerInfo.lastUpdate, time.Now().UnixNano())
	if channel != "" {
		key := ChannelReg{
			PartitionID: pid,
			PeerId:      client.peerInfo.Id,
			Channel:     channel,
		}
		if p.ctx.nsqlookupd.DB.AddChannelReg(topic, key) {
			nsqlookupLog.Logf("DB: client(%s) REGISTER new channel: topic:%s channel:%s pid:%s",
				client, topic, channel, pid)
		}
	}
	if p.ctx.nsqlookupd.DB.AddTopicProducer(topic, pid, &Producer{peerInfo: client.peerInfo}) {
		nsqlookupLog.Logf("DB: client(%s) REGISTER new topic:%s pid:%s",
			client, topic, pid)
	}

	return []byte("OK"), nil
}

func (p *LookupProtocolV1) UNREGISTER(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	if client.peerInfo == nil {
		return nil, protocol.NewFatalClientErr(nil, "E_INVALID", "client must IDENTIFY")
	}

	topic, channel, pid, err := getTopicChan("UNREGISTER", params)
	if err != nil {
		// check if old nsqd
		if client.peerInfo.IsOldPeer() {
			nsqlookupLog.Logf("client %v is old node trying unregister", client)
			pid = strconv.Itoa(OLD_VERSION_PID)
			topic, channel, err = getTopicChanFromOld("UNREGISTER", params)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}

	atomic.StoreInt64(&client.peerInfo.lastUpdate, time.Now().UnixNano())
	if channel != "" {
		key := ChannelReg{
			PartitionID: pid,
			PeerId:      client.peerInfo.Id,
			Channel:     channel,
		}

		removed := p.ctx.nsqlookupd.DB.RemoveChannelReg(topic, key)
		if removed {
			nsqlookupLog.Logf("DB: client(%s) UNREGISTER channel %v on topic:%s-%v",
				client, channel, topic, pid)
		}
	} else {
		// no channel was specified so this is a topic unregistration
		// remove all of the channel registrations...
		// normally this shouldn't happen which is why we print a warning message
		// if anything is actually removed
		key := ChannelReg{
			PartitionID: pid,
			PeerId:      client.peerInfo.Id,
			Channel:     "*",
		}

		removed := p.ctx.nsqlookupd.DB.RemoveChannelReg(topic, key)
		if removed {
			nsqlookupLog.LogWarningf(" client(%s) unexpected UNREGISTER all channels under topic:%s pid:%s",
				client, topic, pid)
		}

		if removed := p.ctx.nsqlookupd.DB.RemoveTopicProducer(topic, pid, client.peerInfo.Id); removed {
			nsqlookupLog.Logf("DB: client(%s) UNREGISTER topic :%s pid:%s",
				client, topic, pid)
		}
	}

	return []byte("OK"), nil
}

func (p *LookupProtocolV1) IDENTIFY(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	var err error

	if client.peerInfo != nil {
		return nil, protocol.NewFatalClientErr(err, "E_INVALID", "cannot IDENTIFY again")
	}

	var bodyLen int32
	err = binary.Read(reader, binary.BigEndian, &bodyLen)
	if err != nil {
		return nil, protocol.NewFatalClientErr(err, "E_BAD_BODY", "IDENTIFY failed to read body size")
	}

	body := make([]byte, bodyLen)
	_, err = io.ReadFull(reader, body)
	if err != nil {
		return nil, protocol.NewFatalClientErr(err, "E_BAD_BODY", "IDENTIFY failed to read body")
	}

	// body is a json structure with producer information
	// Id should be identified by the client.
	peerInfo := PeerInfo{Id: ""}
	err = json.Unmarshal(body, &peerInfo)
	if err != nil {
		return nil, protocol.NewFatalClientErr(err, "E_BAD_BODY", "IDENTIFY failed to decode JSON body")
	}

	peerInfo.RemoteAddress = client.RemoteAddr().String()
	peerInfo.Id = peerInfo.RemoteAddress

	// require all fields
	if peerInfo.Id == "" || peerInfo.BroadcastAddress == "" || peerInfo.TCPPort == 0 || peerInfo.HTTPPort == 0 || peerInfo.Version == "" {
		nsqlookupLog.Logf("identify info missing: %v", peerInfo)
		return nil, protocol.NewFatalClientErr(nil, "E_BAD_BODY", "IDENTIFY missing fields")
	}

	if p.ctx.nsqlookupd.RealTCPAddr() == nil {
		nsqlookupLog.Logf("client(%s) register before the server is ready", client)
		return nil, protocol.NewFatalClientErr(nil, "E_NOT_READY", "The server is not ready for use")
	}

	atomic.StoreInt64(&peerInfo.lastUpdate, time.Now().UnixNano())

	nsqlookupLog.Logf("CLIENT(%s): IDENTIFY peer: %v ",
		client, peerInfo)

	client.peerInfo = &peerInfo
	if p.ctx.nsqlookupd.DB.addPeerClient(peerInfo.Id, client.peerInfo) {
		nsqlookupLog.Logf("DB: client(%s) REGISTER new peer", client)
	}

	// build a response
	data := make(map[string]interface{})
	data["tcp_port"] = p.ctx.nsqlookupd.RealTCPAddr().Port
	_, httpPortStr, _ := net.SplitHostPort(p.ctx.nsqlookupd.opts.HTTPAddress)
	httpPort, _ := strconv.Atoi(httpPortStr)
	if p.ctx.nsqlookupd.RealHTTPAddr() != nil {
		httpPort = p.ctx.nsqlookupd.RealHTTPAddr().Port
	}
	data["http_port"] = httpPort
	data["version"] = version.Binary
	hostname, err := os.Hostname()
	if err != nil {
		nsqlookupLog.LogErrorf("ERROR: unable to get hostname %s", err)
	}
	data["broadcast_address"] = p.ctx.nsqlookupd.opts.BroadcastAddress
	data["hostname"] = hostname

	response, err := json.Marshal(data)
	if err != nil {
		nsqlookupLog.LogErrorf(" marshaling %v", data)
		return []byte("OK"), nil
	}
	return response, nil
}

func (p *LookupProtocolV1) PING(client *ClientV1, params []string) ([]byte, error) {
	if client.peerInfo != nil {
		// we could get a PING before other commands on the same client connection
		cur := time.Unix(0, atomic.LoadInt64(&client.peerInfo.lastUpdate))
		now := time.Now()
		nsqlookupLog.LogDebugf("CLIENT(%s): pinged (last ping %s)", client.peerInfo.Id,
			now.Sub(cur))
		atomic.StoreInt64(&client.peerInfo.lastUpdate, now.UnixNano())
	}
	return []byte("OK"), nil
}
<|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.

// +build freebsd dragonfly darwin linux netbsd openbsd

// This file tests that some basic syscalls are consistent across
// all Unixes.

package syscall_test

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

// {Set,Get}priority and needed constants for them
func _() {
	var (
		_ func(int, int, int) error   = syscall.Setpriority
		_ func(int, int) (int, error) = syscall.Getpriority
	)
	const (
		_ int = syscall.PRIO_USER
		_ int = syscall.PRIO_PROCESS
		_ int = syscall.PRIO_PGRP
	)
}

// termios functions and constants
func _() {
	const (
		_ int = syscall.TCIFLUSH
		_ int = syscall.TCIOFLUSH
		_ int = syscall.TCOFLUSH
	)
}

func _() {
	_ = syscall.Flock_t{
		Type:   int16(0),
		Whence: int16(0),
		Start:  int64(0),
		Len:    int64(0),
		Pid:    int32(0),
	}
}

// TestPassFD tests passing a file descriptor over a Unix socket.
//
// This test involved both a parent and child process. The parent
// process is invoked as a normal test, with "go test", which then
// runs the child process by running the current test binary with args
// "-test.run=^TestPassFD$" and an environment variable used to signal
// that the test should become the child process instead.
func TestPassFD(t *testing.T) {
	if runtime.GOOS == "dragonfly" {
		// TODO(jsing): Figure out why sendmsg is returning EINVAL.
		t.Skip("skipping test on dragonfly")
	}
	if os.Getenv("GO_WANT_HELPER_PROCESS") == "1" {
		passFDChild()
		return
	}

	tempDir, err := ioutil.TempDir("", "TestPassFD")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tempDir)

	fds, err := syscall.Socketpair(syscall.AF_LOCAL, syscall.SOCK_STREAM, 0)
	if err != nil {
		t.Fatalf("Socketpair: %v", err)
	}
	defer syscall.Close(fds[0])
	defer syscall.Close(fds[1])
	writeFile := os.NewFile(uintptr(fds[0]), "child-writes")
	readFile := os.NewFile(uintptr(fds[1]), "parent-reads")
	defer writeFile.Close()
	defer readFile.Close()

	cmd := exec.Command(os.Args[0], "-test.run=^TestPassFD$", "--", tempDir)
	cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1"}
	cmd.ExtraFiles = []*os.File{writeFile}

	out, err := cmd.CombinedOutput()
	if len(out) > 0 || err != nil {
		t.Fatalf("child process: %q, %v", out, err)
	}

	c, err := net.FileConn(readFile)
	if err != nil {
		t.Fatalf("FileConn: %v", err)
	}
	defer c.Close()

	uc, ok := c.(*net.UnixConn)
	if !ok {
		t.Fatalf("unexpected FileConn type; expected UnixConn, got %T", c)
	}

	buf := make([]byte, 32) // expect 1 byte
	oob := make([]byte, 32) // expect 24 bytes
	closeUnix := time.AfterFunc(5*time.Second, func() {
		t.Logf("timeout reading from unix socket")
		uc.Close()
	})
	_, oobn, _, _, err := uc.ReadMsgUnix(buf, oob)
	closeUnix.Stop()

	scms, err := syscall.ParseSocketControlMessage(oob[:oobn])
	if err != nil {
		t.Fatalf("ParseSocketControlMessage: %v", err)
	}
	if len(scms) != 1 {
		t.Fatalf("expected 1 SocketControlMessage; got scms = %#v", scms)
	}
	scm := scms[0]
	gotFds, err := syscall.ParseUnixRights(&scm)
	if err != nil {
		t.Fatalf("syscall.ParseUnixRights: %v", err)
	}
	if len(gotFds) != 1 {
		t.Fatalf("wanted 1 fd; got %#v", gotFds)
	}

	f := os.NewFile(uintptr(gotFds[0]), "fd-from-child")
	defer f.Close()

	got, err := ioutil.ReadAll(f)
	want := "Hello from child process!\n"
	if string(got) != want {
		t.Errorf("child process ReadAll: %q, %v; want %q", got, err, want)
	}
}

// passFDChild is the child process used by TestPassFD.
func passFDChild() {
	defer os.Exit(0)

	// Look for our fd. It should be fd 3, but we work around an fd leak
	// bug here (http://golang.org/issue/2603) to let it be elsewhere.
	var uc *net.UnixConn
	for fd := uintptr(3); fd <= 10; fd++ {
		f := os.NewFile(fd, "unix-conn")
		var ok bool
		netc, _ := net.FileConn(f)
		uc, ok = netc.(*net.UnixConn)
		if ok {
			break
		}
	}
	if uc == nil {
		fmt.Println("failed to find unix fd")
		return
	}

	// Make a file f to send to our parent process on uc.
	// We make it in tempDir, which our parent will clean up.
	flag.Parse()
	tempDir := flag.Arg(0)
	f, err := ioutil.TempFile(tempDir, "")
	if err != nil {
		fmt.Printf("TempFile: %v", err)
		return
	}

	f.Write([]byte("Hello from child process!\n"))
	f.Seek(0, 0)

	rights := syscall.UnixRights(int(f.Fd()))
	dummyByte := []byte("x")
	n, oobn, err := uc.WriteMsgUnix(dummyByte, rights, nil)
	if err != nil {
		fmt.Printf("WriteMsgUnix: %v", err)
		return
	}
	if n != 1 || oobn != len(rights) {
		fmt.Printf("WriteMsgUnix = %d, %d; want 1, %d", n, oobn, len(rights))
		return
	}
}

// TestUnixRightsRoundtrip tests that UnixRights, ParseSocketControlMessage,
// and ParseUnixRights are able to successfully round-trip lists of file descriptors.
func TestUnixRightsRoundtrip(t *testing.T) {
	testCases := [...][][]int{
		{{42}},
		{{1, 2}},
		{{3, 4, 5}},
		{{}},
		{{1, 2}, {3, 4, 5}, {}, {7}},
	}
	for _, testCase := range testCases {
		b := []byte{}
		var n int
		for _, fds := range testCase {
			// Last assignment to n wins
			n = len(b) + syscall.CmsgLen(4*len(fds))
			b = append(b, syscall.UnixRights(fds...)...)
		}
		// Truncate b
		b = b[:n]

		scms, err := syscall.ParseSocketControlMessage(b)
		if err != nil {
			t.Fatalf("ParseSocketControlMessage: %v", err)
		}
		if len(scms) != len(testCase) {
			t.Fatalf("expected %v SocketControlMessage; got scms = %#v", len(testCase), scms)
		}
		for i, scm := range scms {
			gotFds, err := syscall.ParseUnixRights(&scm)
			if err != nil {
				t.Fatalf("ParseUnixRights: %v", err)
			}
			wantFds := testCase[i]
			if len(gotFds) != len(wantFds) {
				t.Fatalf("expected %v fds, got %#v", len(wantFds), gotFds)
			}
			for j, fd := range gotFds {
				if fd != wantFds[j] {
					t.Fatalf("expected fd %v, got %v", wantFds[j], fd)
				}
			}
		}
	}
}

func TestRlimit(t *testing.T) {
	var rlimit, zero syscall.Rlimit
	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rlimit)
	if err != nil {
		t.Fatalf("Getrlimit: save failed: %v", err)
	}
	if zero == rlimit {
		t.Fatalf("Getrlimit: save failed: got zero value %#v", rlimit)
	}
	set := rlimit
	set.Cur = set.Max - 1
	err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &set)
	if err != nil {
		t.Fatalf("Setrlimit: set failed: %#v %v", set, err)
	}
	var get syscall.Rlimit
	err = syscall.Getrlimit(syscall.RLIMIT_NOFILE, &get)
	if err != nil {
		t.Fatalf("Getrlimit: get failed: %v", err)
	}
	set = rlimit
	set.Cur = set.Max - 1
	if set != get {
		// Seems like Darwin requires some privilege to
		// increase the soft limit of rlimit sandbox, though
		// Setrlimit never reports an error.
		switch runtime.GOOS {
		case "darwin":
		default:
			t.Fatalf("Rlimit: change failed: wanted %#v got %#v", set, get)
		}
	}
	err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rlimit)
	if err != nil {
		t.Fatalf("Setrlimit: restore failed: %#v %v", rlimit, err)
	}
}
<commit_msg>syscall: add fcntl test<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.

// +build freebsd dragonfly darwin linux netbsd openbsd

package syscall_test

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"syscall"
	"testing"
	"time"
)

// Tests that below functions, structures and constants are consistent
// on all Unix-like systems.
func _() {
	// program scheduling priority functions and constants
	var (
		_ func(int, int, int) error   = syscall.Setpriority
		_ func(int, int) (int, error) = syscall.Getpriority
	)
	const (
		_ int = syscall.PRIO_USER
		_ int = syscall.PRIO_PROCESS
		_ int = syscall.PRIO_PGRP
	)

	// termios constants
	const (
		_ int = syscall.TCIFLUSH
		_ int = syscall.TCIOFLUSH
		_ int = syscall.TCOFLUSH
	)

	// fcntl file locking structure and constants
	var (
		_ = syscall.Flock_t{
			Type:   int16(0),
			Whence: int16(0),
			Start:  int64(0),
			Len:    int64(0),
			Pid:    int32(0),
		}
	)
	const (
		_ = syscall.F_GETLK
		_ = syscall.F_SETLK
		_ = syscall.F_SETLKW
	)
}

// TestFcntlFlock tests whether the file locking structure matches
// the calling convention of each kernel.
func TestFcntlFlock(t *testing.T) {
	name := filepath.Join(os.TempDir(), "TestFcntlFlock")
	fd, err := syscall.Open(name, syscall.O_CREAT|syscall.O_RDWR|syscall.O_CLOEXEC, 0)
	if err != nil {
		t.Fatalf("Open failed: %v", err)
	}
	defer syscall.Unlink(name)
	defer syscall.Close(fd)
	flock := syscall.Flock_t{
		Type:  syscall.F_RDLCK,
		Start: 0, Len: 0, Whence: 1,
	}
	if err := syscall.FcntlFlock(uintptr(fd), syscall.F_GETLK, &flock); err != nil {
		t.Fatalf("FcntlFlock failed: %v", err)
	}
}

// TestPassFD tests passing a file descriptor over a Unix socket.
//
// This test involved both a parent and child process. The parent
// process is invoked as a normal test, with "go test", which then
// runs the child process by running the current test binary with args
// "-test.run=^TestPassFD$" and an environment variable used to signal
// that the test should become the child process instead.
func TestPassFD(t *testing.T) {
	if runtime.GOOS == "dragonfly" {
		// TODO(jsing): Figure out why sendmsg is returning EINVAL.
		t.Skip("skipping test on dragonfly")
	}
	if os.Getenv("GO_WANT_HELPER_PROCESS") == "1" {
		passFDChild()
		return
	}

	tempDir, err := ioutil.TempDir("", "TestPassFD")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tempDir)

	fds, err := syscall.Socketpair(syscall.AF_LOCAL, syscall.SOCK_STREAM, 0)
	if err != nil {
		t.Fatalf("Socketpair: %v", err)
	}
	defer syscall.Close(fds[0])
	defer syscall.Close(fds[1])
	writeFile := os.NewFile(uintptr(fds[0]), "child-writes")
	readFile := os.NewFile(uintptr(fds[1]), "parent-reads")
	defer writeFile.Close()
	defer readFile.Close()

	cmd := exec.Command(os.Args[0], "-test.run=^TestPassFD$", "--", tempDir)
	cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1"}
	cmd.ExtraFiles = []*os.File{writeFile}

	out, err := cmd.CombinedOutput()
	if len(out) > 0 || err != nil {
		t.Fatalf("child process: %q, %v", out, err)
	}

	c, err := net.FileConn(readFile)
	if err != nil {
		t.Fatalf("FileConn: %v", err)
	}
	defer c.Close()

	uc, ok := c.(*net.UnixConn)
	if !ok {
		t.Fatalf("unexpected FileConn type; expected UnixConn, got %T", c)
	}

	buf := make([]byte, 32) // expect 1 byte
	oob := make([]byte, 32) // expect 24 bytes
	closeUnix := time.AfterFunc(5*time.Second, func() {
		t.Logf("timeout reading from unix socket")
		uc.Close()
	})
	_, oobn, _, _, err := uc.ReadMsgUnix(buf, oob)
	closeUnix.Stop()

	scms, err := syscall.ParseSocketControlMessage(oob[:oobn])
	if err != nil {
		t.Fatalf("ParseSocketControlMessage: %v", err)
	}
	if len(scms) != 1 {
		t.Fatalf("expected 1 SocketControlMessage; got scms = %#v", scms)
	}
	scm := scms[0]
	gotFds, err := syscall.ParseUnixRights(&scm)
	if err != nil {
		t.Fatalf("syscall.ParseUnixRights: %v", err)
	}
	if len(gotFds) != 1 {
		t.Fatalf("wanted 1 fd; got %#v", gotFds)
	}

	f := os.NewFile(uintptr(gotFds[0]), "fd-from-child")
	defer f.Close()

	got, err := ioutil.ReadAll(f)
	want := "Hello from child process!\n"
	if string(got) != want {
		t.Errorf("child process ReadAll: %q, %v; want %q", got, err, want)
	}
}

// passFDChild is the child process used by TestPassFD.
func passFDChild() {
	defer os.Exit(0)

	// Look for our fd. It should be fd 3, but we work around an fd leak
	// bug here (http://golang.org/issue/2603) to let it be elsewhere.
	var uc *net.UnixConn
	for fd := uintptr(3); fd <= 10; fd++ {
		f := os.NewFile(fd, "unix-conn")
		var ok bool
		netc, _ := net.FileConn(f)
		uc, ok = netc.(*net.UnixConn)
		if ok {
			break
		}
	}
	if uc == nil {
		fmt.Println("failed to find unix fd")
		return
	}

	// Make a file f to send to our parent process on uc.
	// We make it in tempDir, which our parent will clean up.
	flag.Parse()
	tempDir := flag.Arg(0)
	f, err := ioutil.TempFile(tempDir, "")
	if err != nil {
		fmt.Printf("TempFile: %v", err)
		return
	}

	f.Write([]byte("Hello from child process!\n"))
	f.Seek(0, 0)

	rights := syscall.UnixRights(int(f.Fd()))
	dummyByte := []byte("x")
	n, oobn, err := uc.WriteMsgUnix(dummyByte, rights, nil)
	if err != nil {
		fmt.Printf("WriteMsgUnix: %v", err)
		return
	}
	if n != 1 || oobn != len(rights) {
		fmt.Printf("WriteMsgUnix = %d, %d; want 1, %d", n, oobn, len(rights))
		return
	}
}

// TestUnixRightsRoundtrip tests that UnixRights, ParseSocketControlMessage,
// and ParseUnixRights are able to successfully round-trip lists of file descriptors.
func TestUnixRightsRoundtrip(t *testing.T) {
	testCases := [...][][]int{
		{{42}},
		{{1, 2}},
		{{3, 4, 5}},
		{{}},
		{{1, 2}, {3, 4, 5}, {}, {7}},
	}
	for _, testCase := range testCases {
		b := []byte{}
		var n int
		for _, fds := range testCase {
			// Last assignment to n wins
			n = len(b) + syscall.CmsgLen(4*len(fds))
			b = append(b, syscall.UnixRights(fds...)...)
		}
		// Truncate b
		b = b[:n]

		scms, err := syscall.ParseSocketControlMessage(b)
		if err != nil {
			t.Fatalf("ParseSocketControlMessage: %v", err)
		}
		if len(scms) != len(testCase) {
			t.Fatalf("expected %v SocketControlMessage; got scms = %#v", len(testCase), scms)
		}
		for i, scm := range scms {
			gotFds, err := syscall.ParseUnixRights(&scm)
			if err != nil {
				t.Fatalf("ParseUnixRights: %v", err)
			}
			wantFds := testCase[i]
			if len(gotFds) != len(wantFds) {
				t.Fatalf("expected %v fds, got %#v", len(wantFds), gotFds)
			}
			for j, fd := range gotFds {
				if fd != wantFds[j] {
					t.Fatalf("expected fd %v, got %v", wantFds[j], fd)
				}
			}
		}
	}
}

func TestRlimit(t *testing.T) {
	var rlimit, zero syscall.Rlimit
	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rlimit)
	if err != nil {
		t.Fatalf("Getrlimit: save failed: %v", err)
	}
	if zero == rlimit {
		t.Fatalf("Getrlimit: save failed: got zero value %#v", rlimit)
	}
	set := rlimit
	set.Cur = set.Max - 1
	err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &set)
	if err != nil {
		t.Fatalf("Setrlimit: set failed: %#v %v", set, err)
	}
	var get syscall.Rlimit
	err = syscall.Getrlimit(syscall.RLIMIT_NOFILE, &get)
	if err != nil {
		t.Fatalf("Getrlimit: get failed: %v", err)
	}
	set = rlimit
	set.Cur = set.Max - 1
	if set != get {
		// Seems like Darwin requires some privilege to
		// increase the soft limit of rlimit sandbox, though
		// Setrlimit never reports an error.
		switch runtime.GOOS {
		case "darwin":
		default:
			t.Fatalf("Rlimit: change failed: wanted %#v got %#v", set, get)
		}
	}
	err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rlimit)
	if err != nil {
		t.Fatalf("Setrlimit: restore failed: %#v %v", rlimit, err)
	}
}
<|endoftext|>