<commit_before>package couchbase

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
)

// The HTTP Client To Use
var HttpClient = http.DefaultClient

// Auth callback gets the auth username and password for the given
// bucket.
type AuthHandler interface {
	GetCredentials() (string, string)
}

type RestPool struct {
	Name         string `json:"name"`
	StreamingURI string `json:"streamingUri"`
	URI          string `json:"uri"`
}

type Pools struct {
	ComponentsVersion     map[string]string `json:"componentsVersion,omitempty"`
	ImplementationVersion string            `json:"implementationVersion"`
	IsAdmin               bool              `json:"isAdminCreds"`
	UUID                  string            `json:"uuid"`
	Pools                 []RestPool        `json:"pools"`
}

// A computer in a cluster running the couchbase software.
type Node struct {
	ClusterCompatibility int                `json:"clusterCompatibility"`
	ClusterMembership    string             `json:"clusterMembership"`
	CouchAPIBase         string             `json:"couchApiBase"`
	Hostname             string             `json:"hostname"`
	InterestingStats     map[string]float64 `json:"interestingStats,omitempty"`
	MCDMemoryAllocated   float64            `json:"mcdMemoryAllocated"`
	MCDMemoryReserved    float64            `json:"mcdMemoryReserved"`
	MemoryFree           float64            `json:"memoryFree"`
	MemoryTotal          float64            `json:"memoryTotal"`
	OS                   string             `json:"os"`
	Ports                map[string]int     `json:"ports"`
	Status               string             `json:"status"`
	Uptime               int                `json:"uptime,string"`
	Version              string             `json:"version"`
	ThisNode             bool               `json:"thisNode,omitempty"`
}

// A pool of nodes and buckets.
type Pool struct {
	BucketMap map[string]Bucket
	Nodes     []Node

	BucketURL map[string]string `json:"buckets"`

	client Client
}

// An individual bucket.  Herein lives the most useful stuff.
type Bucket struct {
	AuthType            string             `json:"authType"`
	Capabilities        []string           `json:"bucketCapabilities"`
	CapabilitiesVersion string             `json:"bucketCapabilitiesVer"`
	Type                string             `json:"bucketType"`
	Name                string             `json:"name"`
	NodeLocator         string             `json:"nodeLocator"`
	Nodes               []Node             `json:"nodes"`
	Quota               map[string]float64 `json:"quota,omitempty"`
	Replicas            int                `json:"replicaNumber"`
	Password            string             `json:"saslPassword"`
	URI                 string             `json:"uri"`
	StreamingURI        string             `json:"streamingUri"`
	LocalRandomKeyURI   string             `json:"localRandomKeyUri,omitempty"`
	UUID                string             `json:"uuid"`
	DDocs               struct {
		URI string `json:"uri"`
	} `json:"ddocs,omitempty"`
	VBucketServerMap struct {
		HashAlgorithm string   `json:"hashAlgorithm"`
		NumReplicas   int      `json:"numReplicas"`
		ServerList    []string `json:"serverList"`
		VBucketMap    [][]int  `json:"vBucketMap"`
	} `json:"vBucketServerMap"`
	BasicStats  map[string]interface{} `json:"basicStats,omitempty"`
	Controllers map[string]interface{} `json:"controllers,omitempty"`

	pool        *Pool
	connections []*connectionPool
	commonSufix string
	auth        AuthHandler
}

// Get the (sorted) list of memcached node addresses (hostname:port).
func (b Bucket) NodeAddresses() []string {
	rv := make([]string, len(b.VBucketServerMap.ServerList))
	copy(rv, b.VBucketServerMap.ServerList)
	sort.Strings(rv)
	return rv
}

// Get the longest common suffix of all host:port strings in the node list.
func (b Bucket) CommonAddressSuffix() string {
	input := []string{}
	for _, n := range b.Nodes {
		input = append(input, n.Hostname)
	}
	return FindCommonSuffix(input)
}

// The couchbase client gives access to all the things.
type Client struct {
	BaseURL  *url.URL
	Info     Pools
	Statuses [256]uint64
}

func maybeAddAuth(req *http.Request, ah AuthHandler) {
	if ah != nil {
		user, pass := ah.GetCredentials()
		req.Header.Set("Authorization", "Basic "+
			base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}

func (c *Client) parseURLResponse(path string, ah AuthHandler,
	out interface{}) error {

	u := *c.BaseURL
	u.User = nil
	if q := strings.Index(path, "?"); q > 0 {
		u.Path = path[:q]
		u.RawQuery = path[q+1:]
	} else {
		u.Path = path
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}
	maybeAddAuth(req, ah)

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	d := json.NewDecoder(res.Body)
	if err = d.Decode(&out); err != nil {
		return err
	}
	return nil
}

type basicAuth struct {
	u, p string
}

func (b basicAuth) GetCredentials() (string, string) {
	return b.u, b.p
}

// Connect to a couchbase cluster.
func Connect(baseU string) (c Client, err error) {
	c.BaseURL, err = url.Parse(baseU)
	if err != nil {
		return
	}

	return c, c.parseURLResponse("/pools", nil, &c.Info)
}

func (b *Bucket) refresh() (err error) {
	pool := b.pool
	err = pool.client.parseURLResponse(b.URI, b.auth, b)
	if err != nil {
		return err
	}
	b.pool = pool
	for i := range b.connections {
		b.connections[i] = newConnectionPool(
			b.VBucketServerMap.ServerList[i], b.auth, 4)
	}
	return nil
}

func (p *Pool) refresh() (err error) {
	p.BucketMap = make(map[string]Bucket)

	var ah AuthHandler
	if user := p.client.BaseURL.User; user != nil {
		pw, _ := user.Password()
		ah = basicAuth{user.Username(), pw}
	}

	buckets := []Bucket{}
	err = p.client.parseURLResponse(p.BucketURL["uri"], ah, &buckets)
	if err != nil {
		return err
	}
	for _, b := range buckets {
		b.pool = p
		b.auth = p.getDefaultAuth(b.Name)
		b.connections = make([]*connectionPool, len(b.VBucketServerMap.ServerList))

		p.BucketMap[b.Name] = b
	}
	return nil
}

// Get a pool from within the couchbase cluster (usually "default").
func (c *Client) GetPool(name string) (p Pool, err error) {
	var poolURI string
	for _, p := range c.Info.Pools {
		if p.Name == name {
			poolURI = p.URI
		}
	}
	if poolURI == "" {
		return p, errors.New("No pool named " + name)
	}

	err = c.parseURLResponse(poolURI, nil, &p)

	p.client = *c

	p.refresh()
	return
}

// Mark this bucket as no longer needed, closing connections it may have open.
func (b *Bucket) Close() {
	if b.connections != nil {
		for _, c := range b.connections {
			if c != nil {
				c.Close()
			}
		}
		b.connections = nil
	}
}

func bucket_finalizer(b *Bucket) {
	if b.connections != nil {
		log.Printf("Warning: Finalizing a bucket with active connections.")
	}
}

func (p *Pool) getDefaultAuth(name string) AuthHandler {
	var pw string
	if p.client.BaseURL.User != nil {
		pw, _ = p.client.BaseURL.User.Password()
	}
	return &basicAuth{name, pw}
}

// Get a bucket from within this pool.
func (p *Pool) GetBucket(name string) (*Bucket, error) {
	ah := p.getDefaultAuth(name)
	rv := &Bucket{}
	err := p.client.parseURLResponse("/pools/default/buckets/"+name, ah, rv)
	if err != nil {
		return nil, err
	}
	rv.pool = p
	rv.auth = ah
	rv.connections = make([]*connectionPool, len(rv.VBucketServerMap.ServerList))

	runtime.SetFinalizer(rv, bucket_finalizer)
	rv.refresh()
	return rv, nil
}

// Get the pool to which this bucket belongs.
func (b *Bucket) GetPool() *Pool {
	return b.pool
}

// Get the client from which we got this pool.
func (p *Pool) GetClient() *Client {
	return &p.client
}

// Convenience function for getting a named bucket from a URL
func GetBucket(endpoint, poolname, bucketname string) (*Bucket, error) {
	var err error
	client, err := Connect(endpoint)
	if err != nil {
		return nil, err
	}

	pool, err := client.GetPool(poolname)
	if err != nil {
		return nil, err
	}

	return pool.GetBucket(bucketname)
}
<commit_msg>Revert to previous GetBucket state<commit_after>package couchbase

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
)

// The HTTP Client To Use
var HttpClient = http.DefaultClient

// Auth callback gets the auth username and password for the given
// bucket.
type AuthHandler interface {
	GetCredentials() (string, string)
}

type RestPool struct {
	Name         string `json:"name"`
	StreamingURI string `json:"streamingUri"`
	URI          string `json:"uri"`
}

type Pools struct {
	ComponentsVersion     map[string]string `json:"componentsVersion,omitempty"`
	ImplementationVersion string            `json:"implementationVersion"`
	IsAdmin               bool              `json:"isAdminCreds"`
	UUID                  string            `json:"uuid"`
	Pools                 []RestPool        `json:"pools"`
}

// A computer in a cluster running the couchbase software.
type Node struct {
	ClusterCompatibility int                `json:"clusterCompatibility"`
	ClusterMembership    string             `json:"clusterMembership"`
	CouchAPIBase         string             `json:"couchApiBase"`
	Hostname             string             `json:"hostname"`
	InterestingStats     map[string]float64 `json:"interestingStats,omitempty"`
	MCDMemoryAllocated   float64            `json:"mcdMemoryAllocated"`
	MCDMemoryReserved    float64            `json:"mcdMemoryReserved"`
	MemoryFree           float64            `json:"memoryFree"`
	MemoryTotal          float64            `json:"memoryTotal"`
	OS                   string             `json:"os"`
	Ports                map[string]int     `json:"ports"`
	Status               string             `json:"status"`
	Uptime               int                `json:"uptime,string"`
	Version              string             `json:"version"`
	ThisNode             bool               `json:"thisNode,omitempty"`
}

// A pool of nodes and buckets.
type Pool struct {
	BucketMap map[string]Bucket
	Nodes     []Node

	BucketURL map[string]string `json:"buckets"`

	client Client
}

// An individual bucket.  Herein lives the most useful stuff.
type Bucket struct {
	AuthType            string             `json:"authType"`
	Capabilities        []string           `json:"bucketCapabilities"`
	CapabilitiesVersion string             `json:"bucketCapabilitiesVer"`
	Type                string             `json:"bucketType"`
	Name                string             `json:"name"`
	NodeLocator         string             `json:"nodeLocator"`
	Nodes               []Node             `json:"nodes"`
	Quota               map[string]float64 `json:"quota,omitempty"`
	Replicas            int                `json:"replicaNumber"`
	Password            string             `json:"saslPassword"`
	URI                 string             `json:"uri"`
	StreamingURI        string             `json:"streamingUri"`
	LocalRandomKeyURI   string             `json:"localRandomKeyUri,omitempty"`
	UUID                string             `json:"uuid"`
	DDocs               struct {
		URI string `json:"uri"`
	} `json:"ddocs,omitempty"`
	VBucketServerMap struct {
		HashAlgorithm string   `json:"hashAlgorithm"`
		NumReplicas   int      `json:"numReplicas"`
		ServerList    []string `json:"serverList"`
		VBucketMap    [][]int  `json:"vBucketMap"`
	} `json:"vBucketServerMap"`
	BasicStats  map[string]interface{} `json:"basicStats,omitempty"`
	Controllers map[string]interface{} `json:"controllers,omitempty"`

	pool        *Pool
	connections []*connectionPool
	commonSufix string
	auth        AuthHandler
}

// Get the (sorted) list of memcached node addresses (hostname:port).
func (b Bucket) NodeAddresses() []string {
	rv := make([]string, len(b.VBucketServerMap.ServerList))
	copy(rv, b.VBucketServerMap.ServerList)
	sort.Strings(rv)
	return rv
}

// Get the longest common suffix of all host:port strings in the node list.
func (b Bucket) CommonAddressSuffix() string {
	input := []string{}
	for _, n := range b.Nodes {
		input = append(input, n.Hostname)
	}
	return FindCommonSuffix(input)
}

// The couchbase client gives access to all the things.
type Client struct {
	BaseURL  *url.URL
	Info     Pools
	Statuses [256]uint64
}

func maybeAddAuth(req *http.Request, ah AuthHandler) {
	if ah != nil {
		user, pass := ah.GetCredentials()
		req.Header.Set("Authorization", "Basic "+
			base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}

func (c *Client) parseURLResponse(path string, ah AuthHandler,
	out interface{}) error {

	u := *c.BaseURL
	u.User = nil
	if q := strings.Index(path, "?"); q > 0 {
		u.Path = path[:q]
		u.RawQuery = path[q+1:]
	} else {
		u.Path = path
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}
	maybeAddAuth(req, ah)

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	d := json.NewDecoder(res.Body)
	if err = d.Decode(&out); err != nil {
		return err
	}
	return nil
}

type basicAuth struct {
	u, p string
}

func (b basicAuth) GetCredentials() (string, string) {
	return b.u, b.p
}

// Connect to a couchbase cluster.
func Connect(baseU string) (c Client, err error) {
	c.BaseURL, err = url.Parse(baseU)
	if err != nil {
		return
	}

	return c, c.parseURLResponse("/pools", nil, &c.Info)
}

func (b *Bucket) refresh() (err error) {
	pool := b.pool
	err = pool.client.parseURLResponse(b.URI, b.auth, b)
	if err != nil {
		return err
	}
	b.pool = pool
	for i := range b.connections {
		b.connections[i] = newConnectionPool(
			b.VBucketServerMap.ServerList[i], b.auth, 4)
	}
	return nil
}

func (p *Pool) refresh() (err error) {
	p.BucketMap = make(map[string]Bucket)

	var ah AuthHandler
	if user := p.client.BaseURL.User; user != nil {
		pw, _ := user.Password()
		ah = basicAuth{user.Username(), pw}
	}

	buckets := []Bucket{}
	err = p.client.parseURLResponse(p.BucketURL["uri"], ah, &buckets)
	if err != nil {
		return err
	}
	for _, b := range buckets {
		b.pool = p
		b.auth = p.getDefaultAuth(b.Name)
		b.connections = make([]*connectionPool, len(b.VBucketServerMap.ServerList))

		p.BucketMap[b.Name] = b
	}
	return nil
}

// Get a pool from within the couchbase cluster (usually "default").
func (c *Client) GetPool(name string) (p Pool, err error) {
	var poolURI string
	for _, p := range c.Info.Pools {
		if p.Name == name {
			poolURI = p.URI
		}
	}
	if poolURI == "" {
		return p, errors.New("No pool named " + name)
	}

	err = c.parseURLResponse(poolURI, nil, &p)

	p.client = *c

	p.refresh()
	return
}

// Mark this bucket as no longer needed, closing connections it may have open.
func (b *Bucket) Close() {
	if b.connections != nil {
		for _, c := range b.connections {
			if c != nil {
				c.Close()
			}
		}
		b.connections = nil
	}
}

func bucket_finalizer(b *Bucket) {
	if b.connections != nil {
		log.Printf("Warning: Finalizing a bucket with active connections.")
	}
}

func (p *Pool) getDefaultAuth(name string) AuthHandler {
	var pw string
	if p.client.BaseURL.User != nil {
		pw, _ = p.client.BaseURL.User.Password()
	}
	return &basicAuth{name, pw}
}

// Get a bucket from within this pool.
func (p *Pool) GetBucket(name string) (*Bucket, error) {
	rv, ok := p.BucketMap[name]
	if !ok {
		return nil, errors.New("No bucket named " + name)
	}
	runtime.SetFinalizer(&rv, bucket_finalizer)
	rv.refresh()
	rv.auth = p.getDefaultAuth(name)
	return &rv, nil
}

// Get the pool to which this bucket belongs.
func (b *Bucket) GetPool() *Pool {
	return b.pool
}

// Get the client from which we got this pool.
func (p *Pool) GetClient() *Client {
	return &p.client
}

// Convenience function for getting a named bucket from a URL
func GetBucket(endpoint, poolname, bucketname string) (*Bucket, error) {
	var err error
	client, err := Connect(endpoint)
	if err != nil {
		return nil, err
	}

	pool, err := client.GetPool(poolname)
	if err != nil {
		return nil, err
	}

	return pool.GetBucket(bucketname)
}
<|endoftext|>


<commit_before>package couchbase

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
)

// The HTTP Client To Use
var HttpClient = http.DefaultClient

// Size of the connection pools (per host).
var PoolSize = 4

// Number of overflow connections allowed in a pool.
var PoolOverflow = PoolSize

// Auth callback gets the auth username and password for the given
// bucket.
type AuthHandler interface {
	GetCredentials() (string, string)
}

type RestPool struct {
	Name         string `json:"name"`
	StreamingURI string `json:"streamingUri"`
	URI          string `json:"uri"`
}

type Pools struct {
	ComponentsVersion     map[string]string `json:"componentsVersion,omitempty"`
	ImplementationVersion string            `json:"implementationVersion"`
	IsAdmin               bool              `json:"isAdminCreds"`
	UUID                  string            `json:"uuid"`
	Pools                 []RestPool        `json:"pools"`
}

// A computer in a cluster running the couchbase software.
type Node struct {
	ClusterCompatibility int                `json:"clusterCompatibility"`
	ClusterMembership    string             `json:"clusterMembership"`
	CouchAPIBase         string             `json:"couchApiBase"`
	Hostname             string             `json:"hostname"`
	InterestingStats     map[string]float64 `json:"interestingStats,omitempty"`
	MCDMemoryAllocated   float64            `json:"mcdMemoryAllocated"`
	MCDMemoryReserved    float64            `json:"mcdMemoryReserved"`
	MemoryFree           float64            `json:"memoryFree"`
	MemoryTotal          float64            `json:"memoryTotal"`
	OS                   string             `json:"os"`
	Ports                map[string]int     `json:"ports"`
	Status               string             `json:"status"`
	Uptime               int                `json:"uptime,string"`
	Version              string             `json:"version"`
	ThisNode             bool               `json:"thisNode,omitempty"`
}

// A pool of nodes and buckets.
type Pool struct {
	BucketMap map[string]Bucket
	Nodes     []Node

	BucketURL map[string]string `json:"buckets"`

	client Client
}

// An individual bucket.  Herein lives the most useful stuff.
type Bucket struct {
	AuthType            string             `json:"authType"`
	Capabilities        []string           `json:"bucketCapabilities"`
	CapabilitiesVersion string             `json:"bucketCapabilitiesVer"`
	Type                string             `json:"bucketType"`
	Name                string             `json:"name"`
	NodeLocator         string             `json:"nodeLocator"`
	Nodes               []Node             `json:"nodes"`
	Quota               map[string]float64 `json:"quota,omitempty"`
	Replicas            int                `json:"replicaNumber"`
	Password            string             `json:"saslPassword"`
	URI                 string             `json:"uri"`
	StreamingURI        string             `json:"streamingUri"`
	LocalRandomKeyURI   string             `json:"localRandomKeyUri,omitempty"`
	UUID                string             `json:"uuid"`
	DDocs               struct {
		URI string `json:"uri"`
	} `json:"ddocs,omitempty"`
	VBucketServerMap struct {
		HashAlgorithm string   `json:"hashAlgorithm"`
		NumReplicas   int      `json:"numReplicas"`
		ServerList    []string `json:"serverList"`
		VBucketMap    [][]int  `json:"vBucketMap"`
	} `json:"vBucketServerMap"`
	BasicStats  map[string]interface{} `json:"basicStats,omitempty"`
	Controllers map[string]interface{} `json:"controllers,omitempty"`

	pool        *Pool
	connPools   []*connectionPool
	commonSufix string
}

func (b Bucket) getConnPools() []*connectionPool {
	return b.connPools
}

func (b Bucket) getConnPool(i int) *connectionPool {
	p := b.getConnPools()
	if len(p) > i {
		return p[i]
	}
	return nil
}

func (b Bucket) authHandler() (ah AuthHandler) {
	if b.pool != nil {
		ah = b.pool.client.ah
	}
	if ah == nil {
		ah = &basicAuth{b.Name, ""}
	}
	return
}

// Get the (sorted) list of memcached node addresses (hostname:port).
func (b Bucket) NodeAddresses() []string {
	rv := make([]string, len(b.VBucketServerMap.ServerList))
	copy(rv, b.VBucketServerMap.ServerList)
	sort.Strings(rv)
	return rv
}

// Get the longest common suffix of all host:port strings in the node list.
func (b Bucket) CommonAddressSuffix() string {
	input := []string{}
	for _, n := range b.Nodes {
		input = append(input, n.Hostname)
	}
	return FindCommonSuffix(input)
}

// The couchbase client gives access to all the things.
type Client struct {
	BaseURL  *url.URL
	ah       AuthHandler
	Info     Pools
	Statuses [256]uint64
}

func maybeAddAuth(req *http.Request, ah AuthHandler) {
	if ah != nil {
		user, pass := ah.GetCredentials()
		req.Header.Set("Authorization", "Basic "+
			base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}

func (c *Client) parseURLResponse(path string, out interface{}) error {
	u := *c.BaseURL
	u.User = nil
	if q := strings.Index(path, "?"); q > 0 {
		u.Path = path[:q]
		u.RawQuery = path[q+1:]
	} else {
		u.Path = path
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}
	maybeAddAuth(req, c.ah)

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		bod, _ := ioutil.ReadAll(io.LimitReader(res.Body, 512))
		return fmt.Errorf("HTTP error %v getting %q: %s",
			res.Status, u.String(), bod)
	}

	d := json.NewDecoder(res.Body)
	if err = d.Decode(&out); err != nil {
		return err
	}
	return nil
}

type basicAuth struct {
	u, p string
}

func (b basicAuth) GetCredentials() (string, string) {
	return b.u, b.p
}

func basicAuthFromURL(us string) (ah AuthHandler) {
	u, err := ParseURL(us)
	if err != nil {
		return
	}
	if user := u.User; user != nil {
		pw, _ := user.Password()
		ah = basicAuth{user.Username(), pw}
	}
	return
}

// ConnectWithAuth connects to a couchbase cluster with the given
// authentication handler.
func ConnectWithAuth(baseU string, ah AuthHandler) (c Client, err error) {
	c.BaseURL, err = ParseURL(baseU)
	if err != nil {
		return
	}
	c.ah = ah

	return c, c.parseURLResponse("/pools", &c.Info)
}

// Connect to a couchbase cluster.  An authentication handler will be
// created from the userinfo in the URL if provided.
func Connect(baseU string) (Client, error) {
	return ConnectWithAuth(baseU, basicAuthFromURL(baseU))
}

func (b *Bucket) refresh() (err error) {
	pool := b.pool
	err = pool.client.parseURLResponse(b.URI, b)
	if err != nil {
		return err
	}
	b.pool = pool
	for i := range b.connPools {
		b.connPools[i] = newConnectionPool(
			b.VBucketServerMap.ServerList[i],
			b.authHandler(), PoolSize, PoolOverflow)
	}
	return nil
}

func (p *Pool) refresh() (err error) {
	p.BucketMap = make(map[string]Bucket)

	buckets := []Bucket{}
	err = p.client.parseURLResponse(p.BucketURL["uri"], &buckets)
	if err != nil {
		return err
	}
	for _, b := range buckets {
		b.pool = p
		b.connPools = make([]*connectionPool, len(b.VBucketServerMap.ServerList))

		p.BucketMap[b.Name] = b
	}
	return nil
}

// Get a pool from within the couchbase cluster (usually "default").
func (c *Client) GetPool(name string) (p Pool, err error) {
	var poolURI string
	for _, p := range c.Info.Pools {
		if p.Name == name {
			poolURI = p.URI
		}
	}
	if poolURI == "" {
		return p, errors.New("No pool named " + name)
	}

	err = c.parseURLResponse(poolURI, &p)

	p.client = *c

	err = p.refresh()
	return
}

// Mark this bucket as no longer needed, closing connections it may have open.
func (b *Bucket) Close() {
	if b.connPools != nil {
		for _, c := range b.connPools {
			if c != nil {
				c.Close()
			}
		}
		b.connPools = nil
	}
}

func bucket_finalizer(b *Bucket) {
	if b.connPools != nil {
		log.Printf("Warning: Finalizing a bucket with active connections.")
	}
}

// Get a bucket from within this pool.
func (p *Pool) GetBucket(name string) (*Bucket, error) {
	rv, ok := p.BucketMap[name]
	if !ok {
		return nil, errors.New("No bucket named " + name)
	}
	runtime.SetFinalizer(&rv, bucket_finalizer)
	err := rv.refresh()
	if err != nil {
		return nil, err
	}
	return &rv, nil
}

// Get the pool to which this bucket belongs.
func (b *Bucket) GetPool() *Pool {
	return b.pool
}

// Get the client from which we got this pool.
func (p *Pool) GetClient() *Client {
	return &p.client
}

// Convenience function for getting a named bucket from a URL
func GetBucket(endpoint, poolname, bucketname string) (*Bucket, error) {
	var err error
	client, err := Connect(endpoint)
	if err != nil {
		return nil, err
	}

	pool, err := client.GetPool(poolname)
	if err != nil {
		return nil, err
	}

	return pool.GetBucket(bucketname)
}
<commit_msg>Remove named error result from *Bucket.refresh<commit_after>package couchbase

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
)

// The HTTP Client To Use
var HttpClient = http.DefaultClient

// Size of the connection pools (per host).
var PoolSize = 4

// Number of overflow connections allowed in a pool.
var PoolOverflow = PoolSize

// Auth callback gets the auth username and password for the given
// bucket.
type AuthHandler interface {
	GetCredentials() (string, string)
}

type RestPool struct {
	Name         string `json:"name"`
	StreamingURI string `json:"streamingUri"`
	URI          string `json:"uri"`
}

type Pools struct {
	ComponentsVersion     map[string]string `json:"componentsVersion,omitempty"`
	ImplementationVersion string            `json:"implementationVersion"`
	IsAdmin               bool              `json:"isAdminCreds"`
	UUID                  string            `json:"uuid"`
	Pools                 []RestPool        `json:"pools"`
}

// A computer in a cluster running the couchbase software.
type Node struct {
	ClusterCompatibility int                `json:"clusterCompatibility"`
	ClusterMembership    string             `json:"clusterMembership"`
	CouchAPIBase         string             `json:"couchApiBase"`
	Hostname             string             `json:"hostname"`
	InterestingStats     map[string]float64 `json:"interestingStats,omitempty"`
	MCDMemoryAllocated   float64            `json:"mcdMemoryAllocated"`
	MCDMemoryReserved    float64            `json:"mcdMemoryReserved"`
	MemoryFree           float64            `json:"memoryFree"`
	MemoryTotal          float64            `json:"memoryTotal"`
	OS                   string             `json:"os"`
	Ports                map[string]int     `json:"ports"`
	Status               string             `json:"status"`
	Uptime               int                `json:"uptime,string"`
	Version              string             `json:"version"`
	ThisNode             bool               `json:"thisNode,omitempty"`
}

// A pool of nodes and buckets.
type Pool struct {
	BucketMap map[string]Bucket
	Nodes     []Node

	BucketURL map[string]string `json:"buckets"`

	client Client
}

// An individual bucket.  Herein lives the most useful stuff.
type Bucket struct {
	AuthType            string             `json:"authType"`
	Capabilities        []string           `json:"bucketCapabilities"`
	CapabilitiesVersion string             `json:"bucketCapabilitiesVer"`
	Type                string             `json:"bucketType"`
	Name                string             `json:"name"`
	NodeLocator         string             `json:"nodeLocator"`
	Nodes               []Node             `json:"nodes"`
	Quota               map[string]float64 `json:"quota,omitempty"`
	Replicas            int                `json:"replicaNumber"`
	Password            string             `json:"saslPassword"`
	URI                 string             `json:"uri"`
	StreamingURI        string             `json:"streamingUri"`
	LocalRandomKeyURI   string             `json:"localRandomKeyUri,omitempty"`
	UUID                string             `json:"uuid"`
	DDocs               struct {
		URI string `json:"uri"`
	} `json:"ddocs,omitempty"`
	VBucketServerMap struct {
		HashAlgorithm string   `json:"hashAlgorithm"`
		NumReplicas   int      `json:"numReplicas"`
		ServerList    []string `json:"serverList"`
		VBucketMap    [][]int  `json:"vBucketMap"`
	} `json:"vBucketServerMap"`
	BasicStats  map[string]interface{} `json:"basicStats,omitempty"`
	Controllers map[string]interface{} `json:"controllers,omitempty"`

	pool        *Pool
	connPools   []*connectionPool
	commonSufix string
}

func (b Bucket) getConnPools() []*connectionPool {
	return b.connPools
}

func (b Bucket) getConnPool(i int) *connectionPool {
	p := b.getConnPools()
	if len(p) > i {
		return p[i]
	}
	return nil
}

func (b Bucket) authHandler() (ah AuthHandler) {
	if b.pool != nil {
		ah = b.pool.client.ah
	}
	if ah == nil {
		ah = &basicAuth{b.Name, ""}
	}
	return
}

// Get the (sorted) list of memcached node addresses (hostname:port).
func (b Bucket) NodeAddresses() []string {
	rv := make([]string, len(b.VBucketServerMap.ServerList))
	copy(rv, b.VBucketServerMap.ServerList)
	sort.Strings(rv)
	return rv
}

// Get the longest common suffix of all host:port strings in the node list.
func (b Bucket) CommonAddressSuffix() string {
	input := []string{}
	for _, n := range b.Nodes {
		input = append(input, n.Hostname)
	}
	return FindCommonSuffix(input)
}

// The couchbase client gives access to all the things.
type Client struct {
	BaseURL  *url.URL
	ah       AuthHandler
	Info     Pools
	Statuses [256]uint64
}

func maybeAddAuth(req *http.Request, ah AuthHandler) {
	if ah != nil {
		user, pass := ah.GetCredentials()
		req.Header.Set("Authorization", "Basic "+
			base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}

func (c *Client) parseURLResponse(path string, out interface{}) error {
	u := *c.BaseURL
	u.User = nil
	if q := strings.Index(path, "?"); q > 0 {
		u.Path = path[:q]
		u.RawQuery = path[q+1:]
	} else {
		u.Path = path
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}
	maybeAddAuth(req, c.ah)

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		bod, _ := ioutil.ReadAll(io.LimitReader(res.Body, 512))
		return fmt.Errorf("HTTP error %v getting %q: %s",
			res.Status, u.String(), bod)
	}

	d := json.NewDecoder(res.Body)
	if err = d.Decode(&out); err != nil {
		return err
	}
	return nil
}

type basicAuth struct {
	u, p string
}

func (b basicAuth) GetCredentials() (string, string) {
	return b.u, b.p
}

func basicAuthFromURL(us string) (ah AuthHandler) {
	u, err := ParseURL(us)
	if err != nil {
		return
	}
	if user := u.User; user != nil {
		pw, _ := user.Password()
		ah = basicAuth{user.Username(), pw}
	}
	return
}

// ConnectWithAuth connects to a couchbase cluster with the given
// authentication handler.
func ConnectWithAuth(baseU string, ah AuthHandler) (c Client, err error) {
	c.BaseURL, err = ParseURL(baseU)
	if err != nil {
		return
	}
	c.ah = ah

	return c, c.parseURLResponse("/pools", &c.Info)
}

// Connect to a couchbase cluster.  An authentication handler will be
// created from the userinfo in the URL if provided.
func Connect(baseU string) (Client, error) {
	return ConnectWithAuth(baseU, basicAuthFromURL(baseU))
}

func (b *Bucket) refresh() error {
	pool := b.pool
	err := pool.client.parseURLResponse(b.URI, b)
	if err != nil {
		return err
	}
	b.pool = pool
	for i := range b.connPools {
		b.connPools[i] = newConnectionPool(
			b.VBucketServerMap.ServerList[i],
			b.authHandler(), PoolSize, PoolOverflow)
	}
	return nil
}

func (p *Pool) refresh() (err error) {
	p.BucketMap = make(map[string]Bucket)

	buckets := []Bucket{}
	err = p.client.parseURLResponse(p.BucketURL["uri"], &buckets)
	if err != nil {
		return err
	}
	for _, b := range buckets {
		b.pool = p
		b.connPools = make([]*connectionPool, len(b.VBucketServerMap.ServerList))

		p.BucketMap[b.Name] = b
	}
	return nil
}

// Get a pool from within the couchbase cluster (usually "default").
func (c *Client) GetPool(name string) (p Pool, err error) {
	var poolURI string
	for _, p := range c.Info.Pools {
		if p.Name == name {
			poolURI = p.URI
		}
	}
	if poolURI == "" {
		return p, errors.New("No pool named " + name)
	}

	err = c.parseURLResponse(poolURI, &p)

	p.client = *c

	err = p.refresh()
	return
}

// Mark this bucket as no longer needed, closing connections it may have open.
func (b *Bucket) Close() {
	if b.connPools != nil {
		for _, c := range b.connPools {
			if c != nil {
				c.Close()
			}
		}
		b.connPools = nil
	}
}

func bucket_finalizer(b *Bucket) {
	if b.connPools != nil {
		log.Printf("Warning: Finalizing a bucket with active connections.")
	}
}

// Get a bucket from within this pool.
func (p *Pool) GetBucket(name string) (*Bucket, error) {
	rv, ok := p.BucketMap[name]
	if !ok {
		return nil, errors.New("No bucket named " + name)
	}
	runtime.SetFinalizer(&rv, bucket_finalizer)
	err := rv.refresh()
	if err != nil {
		return nil, err
	}
	return &rv, nil
}

// Get the pool to which this bucket belongs.
func (b *Bucket) GetPool() *Pool {
	return b.pool
}

// Get the client from which we got this pool.
func (p *Pool) GetClient() *Client {
	return &p.client
}

// Convenience function for getting a named bucket from a URL
func GetBucket(endpoint, poolname, bucketname string) (*Bucket, error) {
	var err error
	client, err := Connect(endpoint)
	if err != nil {
		return nil, err
	}

	pool, err := client.GetPool(poolname)
	if err != nil {
		return nil, err
	}

	return pool.GetBucket(bucketname)
}
<|endoftext|>


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

package strconv_test

import (
	. "strconv"
	"testing"
)

type itob64Test struct {
	in   int64
	base int
	out  string
}

var itob64tests = []itob64Test{
	{0, 10, "0"},
	{1, 10, "1"},
	{-1, 10, "-1"},
	{12345678, 10, "12345678"},
	{-987654321, 10, "-987654321"},
	{1<<31 - 1, 10, "2147483647"},
	{-1<<31 + 1, 10, "-2147483647"},
	{1 << 31, 10, "2147483648"},
	{-1 << 31, 10, "-2147483648"},
	{1<<31 + 1, 10, "2147483649"},
	{-1<<31 - 1, 10, "-2147483649"},
	{1<<32 - 1, 10, "4294967295"},
	{-1<<32 + 1, 10, "-4294967295"},
	{1 << 32, 10, "4294967296"},
	{-1 << 32, 10, "-4294967296"},
	{1<<32 + 1, 10, "4294967297"},
	{-1<<32 - 1, 10, "-4294967297"},
	{1 << 50, 10, "1125899906842624"},
	{1<<63 - 1, 10, "9223372036854775807"},
	{-1<<63 + 1, 10, "-9223372036854775807"},
	{-1 << 63, 10, "-9223372036854775808"},

	{0, 2, "0"},
	{10, 2, "1010"},
	{-1, 2, "-1"},
	{1 << 15, 2, "1000000000000000"},

	{-8, 8, "-10"},
	{057635436545, 8, "57635436545"},
	{1 << 24, 8, "100000000"},

	{16, 16, "10"},
	{-0x123456789abcdef, 16, "-123456789abcdef"},
	{1<<63 - 1, 16, "7fffffffffffffff"},
	{1<<63 - 1, 2, "111111111111111111111111111111111111111111111111111111111111111"},

	{16, 17, "g"},
	{25, 25, "10"},
	{(((((17*35+24)*35+21)*35+34)*35+12)*35+24)*35 + 32, 35, "holycow"},
	{(((((17*36+24)*36+21)*36+34)*36+12)*36+24)*36 + 32, 36, "holycow"},
}

func TestItoa(t *testing.T) {
	for _, test := range itob64tests {
		s := FormatInt(test.in, test.base)
		if s != test.out {
			t.Errorf("FormatInt(%v, %v) = %v want %v",
				test.in, test.base, s, test.out)
		}
		x := AppendInt([]byte("abc"), test.in, test.base)
		if string(x) != "abc"+test.out {
			t.Errorf("AppendInt(%q, %v, %v) = %q want %v",
				"abc", test.in, test.base, x, test.out)
		}

		if test.in >= 0 {
			s := FormatUint(uint64(test.in), test.base)
			if s != test.out {
				t.Errorf("FormatUint(%v, %v) = %v want %v",
					test.in, test.base, s, test.out)
			}
			x := AppendUint(nil, uint64(test.in), test.base)
			if string(x) != test.out {
				t.Errorf("AppendUint(%q, %v, %v) = %q want %v",
					"abc", uint64(test.in), test.base, x, test.out)
			}
		}

		if test.base == 10 && int64(int(test.in)) == test.in {
			s := Itoa(int(test.in))
			if s != test.out {
				t.Errorf("Itoa(%v) = %v want %v",
					test.in, s, test.out)
			}
		}
	}
}

type uitob64Test struct {
	in   uint64
	base int
	out  string
}

var uitob64tests = []uitob64Test{
	{1<<63 - 1, 10, "9223372036854775807"},
	{1 << 63, 10, "9223372036854775808"},
	{1<<63 + 1, 10, "9223372036854775809"},
	{1<<64 - 2, 10, "18446744073709551614"},
	{1<<64 - 1, 10, "18446744073709551615"},
	{1<<64 - 1, 2, "1111111111111111111111111111111111111111111111111111111111111111"},
}

func TestUitoa(t *testing.T) {
	for _, test := range uitob64tests {
		s := FormatUint(test.in, test.base)
		if s != test.out {
			t.Errorf("FormatUint(%v, %v) = %v want %v",
				test.in, test.base, s, test.out)
		}
		x := AppendUint([]byte("abc"), test.in, test.base)
		if string(x) != "abc"+test.out {
			t.Errorf("AppendUint(%q, %v, %v) = %q want %v",
				"abc", test.in, test.base, x, test.out)
		}

	}
}

func BenchmarkFormatInt(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range itob64tests {
			FormatInt(test.in, test.base)
		}
	}
}

func BenchmarkAppendInt(b *testing.B) {
	dst := make([]byte, 0, 30)
	for i := 0; i < b.N; i++ {
		for _, test := range itob64tests {
			AppendInt(dst, test.in, test.base)
		}
	}
}

func BenchmarkFormatUint(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range uitob64tests {
			FormatUint(test.in, test.base)
		}
	}
}

func BenchmarkAppendUint(b *testing.B) {
	dst := make([]byte, 0, 30)
	for i := 0; i < b.N; i++ {
		for _, test := range uitob64tests {
			AppendUint(dst, test.in, test.base)
		}
	}
}
<commit_msg>strconv/itoa: add test to generate the longest output string possible by formatBits<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strconv_test

import (
	. "strconv"
	"testing"
)

type itob64Test struct {
	in   int64
	base int
	out  string
}

var itob64tests = []itob64Test{
	{0, 10, "0"},
	{1, 10, "1"},
	{-1, 10, "-1"},
	{12345678, 10, "12345678"},
	{-987654321, 10, "-987654321"},
	{1<<31 - 1, 10, "2147483647"},
	{-1<<31 + 1, 10, "-2147483647"},
	{1 << 31, 10, "2147483648"},
	{-1 << 31, 10, "-2147483648"},
	{1<<31 + 1, 10, "2147483649"},
	{-1<<31 - 1, 10, "-2147483649"},
	{1<<32 - 1, 10, "4294967295"},
	{-1<<32 + 1, 10, "-4294967295"},
	{1 << 32, 10, "4294967296"},
	{-1 << 32, 10, "-4294967296"},
	{1<<32 + 1, 10, "4294967297"},
	{-1<<32 - 1, 10, "-4294967297"},
	{1 << 50, 10, "1125899906842624"},
	{1<<63 - 1, 10, "9223372036854775807"},
	{-1<<63 + 1, 10, "-9223372036854775807"},
	{-1 << 63, 10, "-9223372036854775808"},

	{0, 2, "0"},
	{10, 2, "1010"},
	{-1, 2, "-1"},
	{1 << 15, 2, "1000000000000000"},

	{-8, 8, "-10"},
	{057635436545, 8, "57635436545"},
	{1 << 24, 8, "100000000"},

	{16, 16, "10"},
	{-0x123456789abcdef, 16, "-123456789abcdef"},
	{1<<63 - 1, 16, "7fffffffffffffff"},
	{1<<63 - 1, 2, "111111111111111111111111111111111111111111111111111111111111111"},
	{-1 << 63, 2, "-1000000000000000000000000000000000000000000000000000000000000000"},

	{16, 17, "g"},
	{25, 25, "10"},
	{(((((17*35+24)*35+21)*35+34)*35+12)*35+24)*35 + 32, 35, "holycow"},
	{(((((17*36+24)*36+21)*36+34)*36+12)*36+24)*36 + 32, 36, "holycow"},
}

func TestItoa(t *testing.T) {
	for _, test := range itob64tests {
		s := FormatInt(test.in, test.base)
		if s != test.out {
			t.Errorf("FormatInt(%v, %v) = %v want %v",
				test.in, test.base, s, test.out)
		}
		x := AppendInt([]byte("abc"), test.in, test.base)
		if string(x) != "abc"+test.out {
			t.Errorf("AppendInt(%q, %v, %v) = %q want %v",
				"abc", test.in, test.base, x, test.out)
		}

		if test.in >= 0 {
			s := FormatUint(uint64(test.in), test.base)
			if s != test.out {
				t.Errorf("FormatUint(%v, %v) = %v want %v",
					test.in, test.base, s, test.out)
			}
			x := AppendUint(nil, uint64(test.in), test.base)
			if string(x) != test.out {
				t.Errorf("AppendUint(%q, %v, %v) = %q want %v",
					"abc", uint64(test.in), test.base, x, test.out)
			}
		}

		if test.base == 10 && int64(int(test.in)) == test.in {
			s := Itoa(int(test.in))
			if s != test.out {
				t.Errorf("Itoa(%v) = %v want %v",
					test.in, s, test.out)
			}
		}
	}
}

type uitob64Test struct {
	in   uint64
	base int
	out  string
}

var uitob64tests = []uitob64Test{
	{1<<63 - 1, 10, "9223372036854775807"},
	{1 << 63, 10, "9223372036854775808"},
	{1<<63 + 1, 10, "9223372036854775809"},
	{1<<64 - 2, 10, "18446744073709551614"},
	{1<<64 - 1, 10, "18446744073709551615"},
	{1<<64 - 1, 2, "1111111111111111111111111111111111111111111111111111111111111111"},
}

func TestUitoa(t *testing.T) {
	for _, test := range uitob64tests {
		s := FormatUint(test.in, test.base)
		if s != test.out {
			t.Errorf("FormatUint(%v, %v) = %v want %v",
				test.in, test.base, s, test.out)
		}
		x := AppendUint([]byte("abc"), test.in, test.base)
		if string(x) != "abc"+test.out {
			t.Errorf("AppendUint(%q, %v, %v) = %q want %v",
				"abc", test.in, test.base, x, test.out)
		}

	}
}

func BenchmarkFormatInt(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range itob64tests {
			FormatInt(test.in, test.base)
		}
	}
}

func BenchmarkAppendInt(b *testing.B) {
	dst := make([]byte, 0, 30)
	for i := 0; i < b.N; i++ {
		for _, test := range itob64tests {
			AppendInt(dst, test.in, test.base)
		}
	}
}

func BenchmarkFormatUint(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range uitob64tests {
			FormatUint(test.in, test.base)
		}
	}
}

func BenchmarkAppendUint(b *testing.B) {
	dst := make([]byte, 0, 30)
	for i := 0; i < b.N; i++ {
		for _, test := range uitob64tests {
			AppendUint(dst, test.in, test.base)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

import (
	"fmt"
	"io"
	"strings"
)

func Fprint(w io.Writer, n Node) error {
	p := printer{
		w: w,
	}
	if f, ok := n.(File); ok {
		p.comments = f.Comments
	}
	p.node(n)
	return p.err
}

type printer struct {
	w   io.Writer
	err error

	wantSpace bool

	curLine int
	level   int

	comments []Comment

	stack []Node
}

func (p *printer) nestedBinary() bool {
	if len(p.stack) < 3 {
		return false
	}
	_, ok := p.stack[len(p.stack)-3].(BinaryExpr)
	return ok
}

func (p *printer) inArithm() bool {
	for i := len(p.stack) - 1; i >= 0; i-- {
		switch p.stack[i].(type) {
		case ArithmExpr, LetStmt, CStyleCond, CStyleLoop:
			return true
		case Stmt:
			return false
		}
	}
	return false
}

func (p *printer) compactArithm() bool {
	for i := len(p.stack) - 1; i >= 0; i-- {
		switch p.stack[i].(type) {
		case LetStmt:
			return true
		case ParenExpr:
			return false
		}
	}
	return false
}

var (
	// these never want a following space
	contiguousRight = map[Token]bool{
		DOLLPR:  true,
		LPAREN:  true,
		DLPAREN: true,
		CMDIN:   true,
		DOLLDP:  true,
	}
	// these never want a preceding space
	contiguousLeft = map[Token]bool{
		SEMICOLON:  true,
		DSEMICOLON: true,
		RPAREN:     true,
		DRPAREN:    true,
		COMMA:      true,
	}
)

func (p *printer) space(b byte) {
	if p.err != nil {
		return
	}
	_, p.err = p.w.Write([]byte{b})
	p.wantSpace = false
}

func (p *printer) nonSpaced(a ...interface{}) {
	for _, v := range a {
		if p.err != nil {
			break
		}
		switch x := v.(type) {
		case string:
			if len(x) > 0 {
				last := x[len(x)-1]
				p.wantSpace = !space[last]
			}
			_, p.err = io.WriteString(p.w, x)
			p.curLine += strings.Count(x, "\n")
		case Comment:
			p.wantSpace = true
			_, p.err = fmt.Fprint(p.w, HASH, x.Text)
		case Token:
			p.wantSpace = !contiguousRight[x]
			_, p.err = fmt.Fprint(p.w, x)
		case Node:
			p.node(x)
		}
	}
}

func (p *printer) spaced(a ...interface{}) {
	for _, v := range a {
		if v == nil {
			continue
		}
		if t, ok := v.(Token); ok && contiguousLeft[t] {
		} else if p.wantSpace {
			p.space(' ')
		}
		p.nonSpaced(v)
	}
}

func (p *printer) indent() {
	for i := 0; i < p.level; i++ {
		p.space('\t')
	}
}

func (p *printer) separate(pos Pos, fallback bool) {
	p.commentsUpTo(pos.Line)
	if p.curLine > 0 && pos.Line > p.curLine {
		p.space('\n')
		if pos.Line > p.curLine+1 {
			// preserve single empty lines
			p.space('\n')
		}
		p.indent()
	} else if fallback {
		p.nonSpaced(SEMICOLON)
	}
	p.curLine = pos.Line
}

func (p *printer) sepSemicolon(v interface{}, pos Pos) {
	p.level++
	p.commentsUpTo(pos.Line)
	p.level--
	p.separate(pos, true)
	p.spaced(v)
}

func (p *printer) sepNewline(v interface{}, pos Pos) {
	p.level++
	p.commentsUpTo(pos.Line)
	p.level--
	p.separate(pos, false)
	p.spaced(v)
}

func (p *printer) commentsUpTo(line int) {
	if len(p.comments) < 1 {
		return
	}
	c := p.comments[0]
	if line > 0 && c.Hash.Line >= line {
		return
	}
	p.separate(c.Hash, false)
	p.spaced(c)
	p.comments = p.comments[1:]
	p.commentsUpTo(line)
}

func (p *printer) node(n Node) {
	p.stack = append(p.stack, n)
	switch x := n.(type) {
	case File:
		p.progStmts(x.Stmts)
		p.commentsUpTo(0)
		p.space('\n')
	case Stmt:
		if x.Negated {
			p.spaced(NOT)
		}
		for _, a := range x.Assigns {
			p.spaced(a)
		}
		p.spaced(x.Node)
		anyNewline := false
		for _, r := range x.Redirs {
			if r.OpPos.Line > p.curLine {
				p.spaced("\\\n")
				if !anyNewline {
					p.level++
					anyNewline = true
				}
				p.indent()
			}
			p.separate(r.OpPos, false)
			p.spaced(r.N)
			p.nonSpaced(r.Op, r.Word)
		}
		for _, r := range x.Redirs {
			if r.Op == SHL || r.Op == DHEREDOC {
				p.space('\n')
				p.curLine++
				p.nonSpaced(r.Hdoc, wordStr(unquote(r.Word)))
			}
		}
		if x.Background {
			p.spaced(AND)
		}
	case Assign:
		if x.Name != nil {
			p.spaced(x.Name)
			if x.Append {
				p.nonSpaced(ADD_ASSIGN)
			} else {
				p.nonSpaced(ASSIGN)
			}
		}
		p.nonSpaced(x.Value)
	case Command:
		p.wordJoin(x.Args, true)
	case Subshell:
		p.spaced(LPAREN)
		if len(x.Stmts) == 0 {
			// avoid conflict with ()
			p.space(' ')
		}
		p.stmtJoin(x.Stmts)
		p.sepNewline(RPAREN, x.Rparen)
	case Block:
		p.spaced(LBRACE)
		p.stmtJoin(x.Stmts)
		p.sepSemicolon(RBRACE, x.Rbrace)
	case IfStmt:
		p.spaced(IF, x.Cond, SEMICOLON, THEN)
		p.curLine = x.Then.Line
		p.stmtJoin(x.ThenStmts)
		for _, el := range x.Elifs {
			p.sepSemicolon(ELIF, el.Elif)
			p.spaced(el.Cond, SEMICOLON, THEN)
			p.stmtJoin(el.ThenStmts)
		}
		if len(x.ElseStmts) > 0 {
			p.sepSemicolon(ELSE, x.Else)
			p.stmtJoin(x.ElseStmts)
		}
		p.sepSemicolon(FI, x.Fi)
	case StmtCond:
		p.stmtJoin(x.Stmts)
	case CStyleCond:
		p.spaced(DLPAREN, x.Cond, DRPAREN)
	case WhileStmt:
		p.spaced(WHILE, x.Cond, SEMICOLON, DO)
		p.curLine = x.Do.Line
		p.stmtJoin(x.DoStmts)
		p.sepSemicolon(DONE, x.Done)
	case UntilStmt:
		p.spaced(UNTIL, x.Cond, SEMICOLON, DO)
		p.curLine = x.Do.Line
		p.stmtJoin(x.DoStmts)
		p.sepSemicolon(DONE, x.Done)
	case ForStmt:
		p.spaced(FOR, x.Cond, SEMICOLON, DO)
		p.curLine = x.Do.Line
		p.stmtJoin(x.DoStmts)
		p.sepSemicolon(DONE, x.Done)
	case WordIter:
		p.spaced(x.Name)
		if len(x.List) > 0 {
			p.spaced(IN)
			p.wordJoin(x.List, false)
		}
	case CStyleLoop:
		p.spaced(DLPAREN, x.Init, SEMICOLON, x.Cond,
			SEMICOLON, x.Post, DRPAREN)
	case UnaryExpr:
		if x.Post {
			p.nonSpaced(x.X, x.Op)
		} else {
			p.nonSpaced(x.Op)
			p.wantSpace = false
			p.nonSpaced(x.X)
		}
	case BinaryExpr:
		switch {
		case p.compactArithm():
			p.nonSpaced(x.X, x.Op, x.Y)
		case p.inArithm():
			p.spaced(x.X, x.Op, x.Y)
		default:
			p.spaced(x.X, x.Op)
			if !p.nestedBinary() {
				p.level++
			}
			p.separate(x.Y.Pos(), false)
			p.nonSpaced(x.Y)
			if !p.nestedBinary() {
				p.level--
			}
		}
	case FuncDecl:
		if x.BashStyle {
			p.spaced(FUNCTION)
		}
		p.spaced(x.Name)
		if !x.BashStyle {
			p.nonSpaced(LPAREN, RPAREN)
		}
		p.spaced(x.Body)
	case Word:
		for _, n := range x.Parts {
			p.nonSpaced(n)
		}
	case Lit:
		p.nonSpaced(x.Value)
	case SglQuoted:
		p.nonSpaced(SQUOTE, x.Value, SQUOTE)
	case Quoted:
		p.nonSpaced(x.Quote)
		for _, n := range x.Parts {
			p.nonSpaced(n)
		}
		p.nonSpaced(quotedStop(x.Quote))
	case CmdSubst:
		if x.Backquotes {
			p.nonSpaced(BQUOTE)
			p.wantSpace = false
		} else {
			p.nonSpaced(DOLLPR)
		}
		p.stmtJoin(x.Stmts)
		if x.Backquotes {
			p.wantSpace = false
			p.sepNewline(BQUOTE, x.Right)
		} else {
			p.sepNewline(RPAREN, x.Right)
		}
	case ParamExp:
		if x.Short {
			p.nonSpaced(DOLLAR, x.Param)
			break
		}
		p.nonSpaced(DOLLBR)
		if x.Length {
			p.nonSpaced(HASH)
		}
		p.nonSpaced(x.Param)
		if x.Ind != nil {
			p.nonSpaced(LBRACK, x.Ind.Word, RBRACK)
		}
		if x.Repl != nil {
			if x.Repl.All {
				p.nonSpaced(QUO)
			}
			p.nonSpaced(QUO, x.Repl.Orig, QUO, x.Repl.With)
		} else if x.Exp != nil {
			p.nonSpaced(x.Exp.Op, x.Exp.Word)
		}
		p.nonSpaced(RBRACE)
	case ArithmExpr:
		p.nonSpaced(DOLLDP, x.X, DRPAREN)
	case ParenExpr:
		p.nonSpaced(LPAREN, x.X, RPAREN)
	case CaseStmt:
		p.spaced(CASE, x.Word, IN)
		for _, pl := range x.List {
			p.separate(wordFirstPos(pl.Patterns), false)
			for i, w := range pl.Patterns {
				if i > 0 {
					p.spaced(OR)
				}
				p.spaced(w)
			}
			p.nonSpaced(RPAREN)
			p.stmtJoin(pl.Stmts)
			p.level++
			p.sepNewline(DSEMICOLON, pl.Dsemi)
			if pl.Dsemi == x.Esac {
				p.curLine--
			}
			p.level--
		}
		if len(x.List) == 0 {
			p.sepSemicolon(ESAC, x.Esac)
		} else {
			p.sepNewline(ESAC, x.Esac)
		}
	case DeclStmt:
		if x.Local {
			p.spaced(LOCAL)
		} else {
			p.spaced(DECLARE)
		}
		for _, w := range x.Opts {
			p.spaced(w)
		}
		for _, a := range x.Assigns {
			p.spaced(a)
		}
	case ArrayExpr:
		p.nonSpaced(LPAREN)
		p.wordJoin(x.List, false)
		p.nonSpaced(RPAREN)
	case CmdInput:
		// avoid conflict with <<
		p.spaced(CMDIN)
		p.stmtJoin(x.Stmts)
		p.nonSpaced(RPAREN)
	case EvalStmt:
		p.spaced(EVAL, x.Stmt)
	case LetStmt:
		p.spaced(LET)
		for _, n := range x.Exprs {
			p.spaced(n)
		}
	}
	p.stack = p.stack[:len(p.stack)-1]
}

func (p *printer) wordJoin(ws []Word, keepNewlines bool) {
	anyNewline := false
	for _, w := range ws {
		if keepNewlines && w.Pos().Line > p.curLine {
			p.spaced("\\\n")
			if !anyNewline {
				p.level++
				anyNewline = true
			}
			p.indent()
		}
		p.spaced(w)
	}
	if anyNewline {
		p.level--
	}
}

func (p *printer) progStmts(stmts []Stmt) {
	for i, s := range stmts {
		p.separate(s.Pos(), i > 0)
		p.node(s)
	}
}

func (p *printer) stmtJoin(stmts []Stmt) {
	p.level++
	for i, s := range stmts {
		p.separate(s.Pos(), i > 0)
		p.node(s)
	}
	p.level--
}
<commit_msg>print: join stmts and separated funcs<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

import (
	"fmt"
	"io"
	"strings"
)

func Fprint(w io.Writer, n Node) error {
	p := printer{
		w: w,
	}
	if f, ok := n.(File); ok {
		p.comments = f.Comments
	}
	p.node(n)
	return p.err
}

type printer struct {
	w   io.Writer
	err error

	wantSpace bool

	curLine int
	level   int

	comments []Comment

	stack []Node
}

func (p *printer) nestedBinary() bool {
	if len(p.stack) < 3 {
		return false
	}
	_, ok := p.stack[len(p.stack)-3].(BinaryExpr)
	return ok
}

func (p *printer) inArithm() bool {
	for i := len(p.stack) - 1; i >= 0; i-- {
		switch p.stack[i].(type) {
		case ArithmExpr, LetStmt, CStyleCond, CStyleLoop:
			return true
		case Stmt:
			return false
		}
	}
	return false
}

func (p *printer) compactArithm() bool {
	for i := len(p.stack) - 1; i >= 0; i-- {
		switch p.stack[i].(type) {
		case LetStmt:
			return true
		case ArithmExpr, ParenExpr:
			return false
		}
	}
	return false
}

var (
	// these never want a following space
	contiguousRight = map[Token]bool{
		DOLLPR:  true,
		LPAREN:  true,
		DLPAREN: true,
		CMDIN:   true,
		DOLLDP:  true,
	}
	// these never want a preceding space
	contiguousLeft = map[Token]bool{
		SEMICOLON:  true,
		DSEMICOLON: true,
		RPAREN:     true,
		DRPAREN:    true,
		COMMA:      true,
	}
)

func (p *printer) space(b byte) {
	if p.err != nil {
		return
	}
	_, p.err = p.w.Write([]byte{b})
	p.wantSpace = false
}

func (p *printer) nonSpaced(a ...interface{}) {
	for _, v := range a {
		if p.err != nil {
			break
		}
		switch x := v.(type) {
		case string:
			if len(x) > 0 {
				last := x[len(x)-1]
				p.wantSpace = !space[last]
			}
			_, p.err = io.WriteString(p.w, x)
			p.curLine += strings.Count(x, "\n")
		case Comment:
			p.wantSpace = true
			_, p.err = fmt.Fprint(p.w, HASH, x.Text)
		case Token:
			p.wantSpace = !contiguousRight[x]
			_, p.err = fmt.Fprint(p.w, x)
		case Node:
			p.node(x)
		}
	}
}

func (p *printer) spaced(a ...interface{}) {
	for _, v := range a {
		if v == nil {
			continue
		}
		if t, ok := v.(Token); ok && contiguousLeft[t] {
		} else if p.wantSpace {
			p.space(' ')
		}
		p.nonSpaced(v)
	}
}

func (p *printer) indent() {
	for i := 0; i < p.level; i++ {
		p.space('\t')
	}
}

func (p *printer) separate(pos Pos, fallback bool) {
	p.commentsUpTo(pos.Line)
	if p.curLine > 0 && pos.Line > p.curLine {
		p.space('\n')
		if pos.Line > p.curLine+1 {
			// preserve single empty lines
			p.space('\n')
		}
		p.indent()
	} else if fallback {
		p.nonSpaced(SEMICOLON)
	}
	p.curLine = pos.Line
}

func (p *printer) separated(v interface{}, pos Pos, fallback bool) {
	p.level++
	p.commentsUpTo(pos.Line)
	p.level--
	p.separate(pos, fallback)
	p.spaced(v)
}

func (p *printer) commentsUpTo(line int) {
	if len(p.comments) < 1 {
		return
	}
	c := p.comments[0]
	if line > 0 && c.Hash.Line >= line {
		return
	}
	p.separate(c.Hash, false)
	p.spaced(c)
	p.comments = p.comments[1:]
	p.commentsUpTo(line)
}

func (p *printer) node(n Node) {
	p.stack = append(p.stack, n)
	switch x := n.(type) {
	case File:
		p.stmts(x.Stmts)
		p.commentsUpTo(0)
		p.space('\n')
	case Stmt:
		if x.Negated {
			p.spaced(NOT)
		}
		for _, a := range x.Assigns {
			p.spaced(a)
		}
		p.spaced(x.Node)
		anyNewline := false
		for _, r := range x.Redirs {
			if r.OpPos.Line > p.curLine {
				p.spaced("\\\n")
				if !anyNewline {
					p.level++
					anyNewline = true
				}
				p.indent()
			}
			p.separate(r.OpPos, false)
			p.spaced(r.N)
			p.nonSpaced(r.Op, r.Word)
		}
		for _, r := range x.Redirs {
			if r.Op == SHL || r.Op == DHEREDOC {
				p.space('\n')
				p.curLine++
				p.nonSpaced(r.Hdoc, wordStr(unquote(r.Word)))
			}
		}
		if x.Background {
			p.spaced(AND)
		}
	case Assign:
		if x.Name != nil {
			p.spaced(x.Name)
			if x.Append {
				p.nonSpaced(ADD_ASSIGN)
			} else {
				p.nonSpaced(ASSIGN)
			}
		}
		p.nonSpaced(x.Value)
	case Command:
		p.wordJoin(x.Args, true)
	case Subshell:
		p.spaced(LPAREN)
		if len(x.Stmts) == 0 {
			// avoid conflict with ()
			p.space(' ')
		}
		p.nestedStmts(x.Stmts)
		p.separated(RPAREN, x.Rparen, false)
	case Block:
		p.spaced(LBRACE)
		p.nestedStmts(x.Stmts)
		p.separated(RBRACE, x.Rbrace, true)
	case IfStmt:
		p.spaced(IF, x.Cond, SEMICOLON, THEN)
		p.curLine = x.Then.Line
		p.nestedStmts(x.ThenStmts)
		for _, el := range x.Elifs {
			p.separated(ELIF, el.Elif, true)
			p.spaced(el.Cond, SEMICOLON, THEN)
			p.nestedStmts(el.ThenStmts)
		}
		if len(x.ElseStmts) > 0 {
			p.separated(ELSE, x.Else, true)
			p.nestedStmts(x.ElseStmts)
		}
		p.separated(FI, x.Fi, true)
	case StmtCond:
		p.nestedStmts(x.Stmts)
	case CStyleCond:
		p.spaced(DLPAREN, x.Cond, DRPAREN)
	case WhileStmt:
		p.spaced(WHILE, x.Cond, SEMICOLON, DO)
		p.curLine = x.Do.Line
		p.nestedStmts(x.DoStmts)
		p.separated(DONE, x.Done, true)
	case UntilStmt:
		p.spaced(UNTIL, x.Cond, SEMICOLON, DO)
		p.curLine = x.Do.Line
		p.nestedStmts(x.DoStmts)
		p.separated(DONE, x.Done, true)
	case ForStmt:
		p.spaced(FOR, x.Cond, SEMICOLON, DO)
		p.curLine = x.Do.Line
		p.nestedStmts(x.DoStmts)
		p.separated(DONE, x.Done, true)
	case WordIter:
		p.spaced(x.Name)
		if len(x.List) > 0 {
			p.spaced(IN)
			p.wordJoin(x.List, false)
		}
	case CStyleLoop:
		p.spaced(DLPAREN, x.Init, SEMICOLON, x.Cond,
			SEMICOLON, x.Post, DRPAREN)
	case UnaryExpr:
		if x.Post {
			p.nonSpaced(x.X, x.Op)
		} else {
			p.nonSpaced(x.Op)
			p.wantSpace = false
			p.nonSpaced(x.X)
		}
	case BinaryExpr:
		switch {
		case p.compactArithm():
			p.nonSpaced(x.X, x.Op, x.Y)
		case p.inArithm():
			p.spaced(x.X, x.Op, x.Y)
		default:
			p.spaced(x.X, x.Op)
			if !p.nestedBinary() {
				p.level++
			}
			p.separate(x.Y.Pos(), false)
			p.nonSpaced(x.Y)
			if !p.nestedBinary() {
				p.level--
			}
		}
	case FuncDecl:
		if x.BashStyle {
			p.spaced(FUNCTION)
		}
		p.spaced(x.Name)
		if !x.BashStyle {
			p.nonSpaced(LPAREN, RPAREN)
		}
		p.spaced(x.Body)
	case Word:
		for _, n := range x.Parts {
			p.nonSpaced(n)
		}
	case Lit:
		p.nonSpaced(x.Value)
	case SglQuoted:
		p.nonSpaced(SQUOTE, x.Value, SQUOTE)
	case Quoted:
		p.nonSpaced(x.Quote)
		for _, n := range x.Parts {
			p.nonSpaced(n)
		}
		p.nonSpaced(quotedStop(x.Quote))
	case CmdSubst:
		if x.Backquotes {
			p.nonSpaced(BQUOTE)
			p.wantSpace = false
		} else {
			p.nonSpaced(DOLLPR)
		}
		p.nestedStmts(x.Stmts)
		if x.Backquotes {
			p.wantSpace = false
			p.separated(BQUOTE, x.Right, false)
		} else {
			p.separated(RPAREN, x.Right, false)
		}
	case ParamExp:
		if x.Short {
			p.nonSpaced(DOLLAR, x.Param)
			break
		}
		p.nonSpaced(DOLLBR)
		if x.Length {
			p.nonSpaced(HASH)
		}
		p.nonSpaced(x.Param)
		if x.Ind != nil {
			p.nonSpaced(LBRACK, x.Ind.Word, RBRACK)
		}
		if x.Repl != nil {
			if x.Repl.All {
				p.nonSpaced(QUO)
			}
			p.nonSpaced(QUO, x.Repl.Orig, QUO, x.Repl.With)
		} else if x.Exp != nil {
			p.nonSpaced(x.Exp.Op, x.Exp.Word)
		}
		p.nonSpaced(RBRACE)
	case ArithmExpr:
		p.nonSpaced(DOLLDP, x.X, DRPAREN)
	case ParenExpr:
		p.nonSpaced(LPAREN, x.X, RPAREN)
	case CaseStmt:
		p.spaced(CASE, x.Word, IN)
		for _, pl := range x.List {
			p.separate(wordFirstPos(pl.Patterns), false)
			for i, w := range pl.Patterns {
				if i > 0 {
					p.spaced(OR)
				}
				p.spaced(w)
			}
			p.nonSpaced(RPAREN)
			p.nestedStmts(pl.Stmts)
			p.level++
			p.separated(DSEMICOLON, pl.Dsemi, false)
			if pl.Dsemi == x.Esac {
				p.curLine--
			}
			p.level--
		}
		p.separated(ESAC, x.Esac, len(x.List) == 0)
	case DeclStmt:
		if x.Local {
			p.spaced(LOCAL)
		} else {
			p.spaced(DECLARE)
		}
		for _, w := range x.Opts {
			p.spaced(w)
		}
		for _, a := range x.Assigns {
			p.spaced(a)
		}
	case ArrayExpr:
		p.nonSpaced(LPAREN)
		p.wordJoin(x.List, false)
		p.nonSpaced(RPAREN)
	case CmdInput:
		// avoid conflict with <<
		p.spaced(CMDIN)
		p.nestedStmts(x.Stmts)
		p.nonSpaced(RPAREN)
	case EvalStmt:
		p.spaced(EVAL, x.Stmt)
	case LetStmt:
		p.spaced(LET)
		for _, n := range x.Exprs {
			p.spaced(n)
		}
	}
	p.stack = p.stack[:len(p.stack)-1]
}

func (p *printer) wordJoin(ws []Word, keepNewlines bool) {
	anyNewline := false
	for _, w := range ws {
		if keepNewlines && w.Pos().Line > p.curLine {
			p.spaced("\\\n")
			if !anyNewline {
				p.level++
				anyNewline = true
			}
			p.indent()
		}
		p.spaced(w)
	}
	if anyNewline {
		p.level--
	}
}

func (p *printer) stmts(stmts []Stmt) {
	for i, s := range stmts {
		p.separate(s.Pos(), i > 0)
		p.node(s)
	}
}

func (p *printer) nestedStmts(stmts []Stmt) {
	p.level++
	p.stmts(stmts)
	p.level--
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"strconv"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"log"
	"net/http"
	"fmt"
	//"regexp"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/leibowitz/goproxy"
	"github.com/twinj/uuid"
)

type ContextUserData struct {
	Store bool
	Time  int64
	Body   io.Reader
	//Body   []byte
	Header http.Header
	Origin string
}

type Content struct {
	//Id       bson.ObjectId
	Request  Request   "request"
	Response Response  "response"
	Date     time.Time "date"
	SocketUUID uuid.UUID "uuid"
}

type Rule struct {
	//Id       bson.ObjectId
	Active	bool "active"
	Dynamic bool "dynamic"
	Host	string "host"
	Path	string "path"
	Query	string "query"
	Method	string "method"
	Status	string "status"
	Response string "response"
	Body	string "body"
	ReqBody string "reqbody"
	ReqHeader http.Header "reqheaders"
	RespHeader http.Header "respheaders"
	Origin  string "origin"
}

type Request struct {
	Origin	string "origin"
	Body   string "body"
	FileId bson.ObjectId
	Query  string "query"
	//Date    time.Time   "date"
	Host    string      "host"
	Method  string      "method"
	Path    string      "path"
	Time    float32     "time"
	Headers http.Header "headers"
}

type Response struct {
	FileId  bson.ObjectId
	Body    string      "body"
	Headers http.Header "headers"
	Status  int         "status"
}

func NewResponse(r *http.Request, headers http.Header, status int, body io.ReadCloser) *http.Response {
	resp := &http.Response{}
	resp.Request = r
	resp.TransferEncoding = r.TransferEncoding
	resp.Header = headers
	resp.StatusCode = status
	//resp.ContentLength = int64(buf.Len())
	resp.Body = body
	return resp
}

func main() {
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	mongourl := flag.String("mongourl", "", "record request/response in mongodb")
	mock := flag.Bool("m", false, "send fake responses")
	addr := flag.String("l", ":8080", "on which address should the proxy listen")

	flag.Parse()

	tmpdir := filepath.Join(os.TempDir(), "proxy-service")

	if _, err := os.Stat(tmpdir); err != nil {
		if os.IsNotExist(err) {
			// Create temp directory to store body response
			err = os.MkdirAll(tmpdir, 0777)
		}

		// err should be nil if we just created the directory
		if err != nil {
			panic(err)
		}
	}

	db := new(mgo.Database)
	c := new(mgo.Collection)
	rules := new(mgo.Collection)

	if len(*mongourl) != 0 {
		// Mongo DB connection
		session, err := mgo.Dial(*mongourl)
		if err != nil {
			panic(err)
		}
		defer session.Close()

		// Optional. Switch the session to a monotonic behavior.
		session.SetMode(mgo.Monotonic, true)

		db = session.DB("proxyservice")
		c = db.C("log_logentry")
		rules = db.C("log_rules")
	}

	uuid.SwitchFormat(uuid.CleanHyphen, false)
	proxy := goproxy.NewProxyHttpServer()
	proxy.Verbose = *verbose

	proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm)

	proxy.OnRequest().DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		origin := ipAddrFromRemoteAddr(ctx.Req.RemoteAddr)

		//log.Printf("Request: %s %s %s", req.Method, req.Host, req.RequestURI)

		host := req.Host
		// Request to domain--name-co-uk.mocky.dev
		// will be forwarded to domain-name.co.uk
		/*if strings.Contains(host, ".mocky.dev") {
			host = strings.Replace(host, ".mocky.dev", "", 1)
			host = strings.Replace(host, "-", ".", -1)
			host = strings.Replace(host, "..", "-", -1)
		} else if strings.Contains(host, ".proxy.dev") {
			host = strings.Replace(host, ".proxy.dev", "", 1)
			host = strings.Replace(host, "-", ".", -1)
			host = strings.Replace(host, "..", "-", -1)
		}*/

		/*r, _ := regexp.Compile(".([0-9]+)$")
		// Check if host is hostname.80 (host with port number)
		res := r.FindStringSubmatch(host)
		if res != nil && len(res[1]) != 0 {
			host = strings.Replace(host, strings.Join([]string{".", res[1]}, ""), "", 1)
			host = strings.Join([]string{host, res[1]}, ":")
			log.Printf("Changing host to %v", host);
		}*/

		//log.Printf("Target Host: %s - Headers: %+v", host, req.Header)
		req.Host = host

		//log.Printf("%+v", req)

		var reqbody []byte

		var bodyreader io.Reader
		if rules.Database != nil && *mock && req.Method != "CONNECT" {
			//reqbody := string(body[:])
			//log.Printf("request body: %s", reqbody)
			rule := Rule{}
			//ctx.Logf("Looking for existing request")
			/*fmt.Println("RequestURI:", req.RequestURI)
			  fmt.Println("Path:", req.URL.Path)
			  fmt.Println("Host:", req.Host)
			  fmt.Println("Method:", req.Method)*/
b := bson.M{"$and": []bson.M{
	bson.M{"active": true},
	bson.M{"dynamic": false},
	bson.M{"origin":
	    bson.M{"$in":
		[]interface{}{origin, false},
	    },
	},
	bson.M{"host":
	    bson.M{"$in":
		[]interface{}{req.Host, false},
	    },
	},
	bson.M{"method":
	    bson.M{"$in":
		[]interface{}{req.Method, false},
	    },
	},
	bson.M{"path":
	    bson.M{"$in":
		[]interface{}{req.URL.Path, false},
	    },
	},
	bson.M{"query":
	    bson.M{"$in":
		[]interface{}{req.URL.Query().Encode(), false},
	    },
	},
    }}

			//b := bson.M{"active": true, "dynamic": false, "host": req.Host, "method": req.Method, "path": req.URL.Path, "query": req.URL.Query().Encode()}
			err := rules.Find(b).One(&rule)//.Sort("priority")
			//log.Printf("Query: %+v, Res: %+v", b, rule)
			if err == nil {
			    status, err := strconv.Atoi(rule.Status)
			    reqbody := ioutil.NopCloser(bytes.NewBufferString(rule.ReqBody))
			    respbody := ioutil.NopCloser(bytes.NewBufferString(rule.Body))
			    log.Printf("%+v", rule)
			    resp := NewResponse(req, rule.RespHeader, status, respbody)
			    ctx.UserData = ContextUserData{Store: true, Time: 0, Body: reqbody, Header: rule.RespHeader, Origin: origin}
			    return req, resp
			}

			// read the whole body
			reqbody, err = ioutil.ReadAll(req.Body)
			if err != nil {
				ctx.Warnf("Cannot read request body %s", err)
			}

			defer req.Body.Close()
			req.Body = ioutil.NopCloser(bytes.NewBuffer(reqbody))

			bodyreader = bytes.NewReader(reqbody)

		} else {
		    bodyreader = req.Body
		}

		ctx.UserData = ContextUserData{Store: true, Time: time.Now().UnixNano(), Body: bodyreader, Header: req.Header, Origin: origin}
		return req, nil
	})

	proxy.OnResponse().DoFunc(func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
		//ctx.Logf("Method: %s - host: %s", ctx.Resp.Request.Method, ctx.Resp.Request.Host)
		if c.Database != nil && ctx.UserData != nil && ctx.UserData.(ContextUserData).Store && ctx.Resp.Request.Method != "CONNECT" {
			// get response content type
			respctype := getContentType(ctx.Resp.Header.Get("Content-Type"))

			//log.Printf("Resp Contenttype %s", respctype)

			respid := bson.NewObjectId()
			//log.Printf("Resp id: %s, host: %s", respid.Hex(), ctx.Resp.Request.Host)

			filename := filepath.Join(tmpdir, respid.Hex())

			//log.Printf("Duplicating Body file id: %s", respid.String())
			resp.Body = NewTeeReadCloser(resp.Body, NewFileStream(filename, *db, respctype, respid))

			reqctype := getContentType(ctx.Resp.Request.Header.Get("Content-Type"))

			//log.Printf("Req Contenttype %s", reqctype)

			if reqctype == "application/x-www-form-urlencoded" {
				//log.Printf("setting req content type to text/plain for saving to mongo")
				reqctype = "text/plain"
			}

			reqid := bson.NewObjectId()
			//log.Printf("Req id: %s, host: %s", reqid.Hex(), ctx.Resp.Request.Host)

			saveFileToMongo(*db, reqid, reqctype, ctx.UserData.(ContextUserData).Body, reqid.Hex())

			// prepare document
			content := Content{
				//Id: docid,
				Request: Request{
					Origin:	 ctx.UserData.(ContextUserData).Origin,
					Path:    ctx.Resp.Request.URL.Path,
					Query:   ctx.Resp.Request.URL.Query().Encode(),
					FileId:  reqid,
					Host:    ctx.Resp.Request.Host,
					Method:  ctx.Resp.Request.Method,
					Time:    float32(time.Now().UnixNano()-ctx.UserData.(ContextUserData).Time) / 1.0e9,
					Headers: ctx.UserData.(ContextUserData).Header},
				Response: Response{
					Status:  ctx.Resp.StatusCode,
					Headers: ctx.Resp.Header,
					FileId:  respid},
				SocketUUID: ctx.Uuid,
				Date: time.Now(),
			}

			err := c.Insert(content)
			if err != nil {
				ctx.Logf("Can't insert document: %v\n", err)
			}

		}
		return resp
	})

	log.Println("Starting Proxy")
	log.Fatalln(http.ListenAndServe(*addr, proxy))
}

type TeeReadCloser struct {
	r io.Reader
	w io.WriteCloser
	c io.Closer
}

func NewTeeReadCloser(r io.ReadCloser, w io.WriteCloser) io.ReadCloser {
	return &TeeReadCloser{io.TeeReader(r, w), w, r}
}

func (t *TeeReadCloser) Read(b []byte) (int, error) {
	return t.r.Read(b)
}

func (t *TeeReadCloser) Close() error {
	err1 := t.c.Close()
	err2 := t.w.Close()
	if err1 == nil && err2 == nil {
		return nil
	}
	if err1 != nil {
		return err2
	}
	return err1
}

type FileStream struct {
	path        string
	db          mgo.Database
	contentType string
	objectId    bson.ObjectId
	f           *os.File
}

func NewFileStream(path string, db mgo.Database, contentType string, objectId bson.ObjectId) *FileStream {
	return &FileStream{path: path, db: db, contentType: contentType, objectId: objectId, f: nil}
}

func (fs *FileStream) Write(b []byte) (nr int, err error) {
	if fs.f == nil {
		fs.f, err = os.Create(fs.path)
		if err != nil {
			return 0, err
		}
	}
	return fs.f.Write(b)
}

func (fs *FileStream) Close() error {
	if fs.f == nil {
		return errors.New("FileStream was never written into")
	}
	fs.f.Seek(0, 0)
	saveFileToMongo(fs.db, fs.objectId, fs.contentType, fs.f, fs.objectId.Hex())
	err := fs.f.Close()
	if err == nil {
		err2 := os.Remove(fs.path)
		if err2 != nil {
			log.Printf("Unable to delete file")
		}
	}
	return err
}

func getMongoFileContent(db mgo.Database, objId bson.ObjectId) (file *mgo.GridFile, err error) {
    file, err = db.GridFS("fs").OpenId(objId)

    if err != nil {
	return file, err
	if err == mgo.ErrNotFound {
	}
    }
    //defer file.Close()

    return file, err
}

// Store file in MongoDB GridFS
func saveFileToMongo(db mgo.Database, objId bson.ObjectId, contentType string, openFile io.Reader, fileName string) {
	mdbfile, err := db.GridFS("fs").Create(fileName)
	if err == nil {
		mdbfile.SetContentType(contentType)
		mdbfile.SetId(objId)
		_, err = io.Copy(mdbfile, openFile)
		if err != nil {
			log.Printf("Unable to copy to mongo: %s - %v", fileName, err)
		}
		err = mdbfile.Close()
		if err != nil {
			log.Printf("Unable to close copy to mongo")
		}
	}
}

func getContentType(s string) string {
	arr := strings.Split(s, ";")
	return arr[0]
}
func ipAddrFromRemoteAddr(s string) string {
        idx := strings.LastIndex(s, ":")
        if idx == -1 {
                return s
        }
        return s[:idx]
}
<commit_msg>Pass ctx for logging purpose<commit_after>package main

import (
	"flag"
	"strconv"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"log"
	"net/http"
	"fmt"
	//"regexp"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/leibowitz/goproxy"
	"github.com/twinj/uuid"
)

type ContextUserData struct {
	Store bool
	Time  int64
	Body   io.Reader
	//Body   []byte
	Header http.Header
	Origin string
}

type Content struct {
	//Id       bson.ObjectId
	Request  Request   "request"
	Response Response  "response"
	Date     time.Time "date"
	SocketUUID uuid.UUID "uuid"
}

type Rule struct {
	//Id       bson.ObjectId
	Active	bool "active"
	Dynamic bool "dynamic"
	Host	string "host"
	Path	string "path"
	Query	string "query"
	Method	string "method"
	Status	string "status"
	Response string "response"
	Body	string "body"
	ReqBody string "reqbody"
	ReqHeader http.Header "reqheaders"
	RespHeader http.Header "respheaders"
	Origin  string "origin"
}

type Request struct {
	Origin	string "origin"
	Body   string "body"
	FileId bson.ObjectId
	Query  string "query"
	//Date    time.Time   "date"
	Host    string      "host"
	Method  string      "method"
	Path    string      "path"
	Time    float32     "time"
	Headers http.Header "headers"
}

type Response struct {
	FileId  bson.ObjectId
	Body    string      "body"
	Headers http.Header "headers"
	Status  int         "status"
}

func NewResponse(r *http.Request, headers http.Header, status int, body io.ReadCloser) *http.Response {
	resp := &http.Response{}
	resp.Request = r
	resp.TransferEncoding = r.TransferEncoding
	resp.Header = headers
	resp.StatusCode = status
	//resp.ContentLength = int64(buf.Len())
	resp.Body = body
	return resp
}

func main() {
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	mongourl := flag.String("mongourl", "", "record request/response in mongodb")
	mock := flag.Bool("m", false, "send fake responses")
	addr := flag.String("l", ":8080", "on which address should the proxy listen")

	flag.Parse()

	tmpdir := filepath.Join(os.TempDir(), "proxy-service")

	if _, err := os.Stat(tmpdir); err != nil {
		if os.IsNotExist(err) {
			// Create temp directory to store body response
			err = os.MkdirAll(tmpdir, 0777)
		}

		// err should be nil if we just created the directory
		if err != nil {
			panic(err)
		}
	}

	db := new(mgo.Database)
	c := new(mgo.Collection)
	rules := new(mgo.Collection)

	if len(*mongourl) != 0 {
		// Mongo DB connection
		session, err := mgo.Dial(*mongourl)
		if err != nil {
			panic(err)
		}
		defer session.Close()

		// Optional. Switch the session to a monotonic behavior.
		session.SetMode(mgo.Monotonic, true)

		db = session.DB("proxyservice")
		c = db.C("log_logentry")
		rules = db.C("log_rules")
	}

	uuid.SwitchFormat(uuid.CleanHyphen, false)
	proxy := goproxy.NewProxyHttpServer()
	proxy.Verbose = *verbose

	proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm)

	proxy.OnRequest().DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		origin := ipAddrFromRemoteAddr(ctx.Req.RemoteAddr)

		//log.Printf("Request: %s %s %s", req.Method, req.Host, req.RequestURI)

		host := req.Host
		// Request to domain--name-co-uk.mocky.dev
		// will be forwarded to domain-name.co.uk
		/*if strings.Contains(host, ".mocky.dev") {
			host = strings.Replace(host, ".mocky.dev", "", 1)
			host = strings.Replace(host, "-", ".", -1)
			host = strings.Replace(host, "..", "-", -1)
		} else if strings.Contains(host, ".proxy.dev") {
			host = strings.Replace(host, ".proxy.dev", "", 1)
			host = strings.Replace(host, "-", ".", -1)
			host = strings.Replace(host, "..", "-", -1)
		}*/

		/*r, _ := regexp.Compile(".([0-9]+)$")
		// Check if host is hostname.80 (host with port number)
		res := r.FindStringSubmatch(host)
		if res != nil && len(res[1]) != 0 {
			host = strings.Replace(host, strings.Join([]string{".", res[1]}, ""), "", 1)
			host = strings.Join([]string{host, res[1]}, ":")
			log.Printf("Changing host to %v", host);
		}*/

		//log.Printf("Target Host: %s - Headers: %+v", host, req.Header)
		req.Host = host

		//log.Printf("%+v", req)

		var reqbody []byte

		var bodyreader io.Reader
		if rules.Database != nil && *mock && req.Method != "CONNECT" {
			//reqbody := string(body[:])
			//log.Printf("request body: %s", reqbody)
			rule := Rule{}
			//ctx.Logf("Looking for existing request")
			/*fmt.Println("RequestURI:", req.RequestURI)
			  fmt.Println("Path:", req.URL.Path)
			  fmt.Println("Host:", req.Host)
			  fmt.Println("Method:", req.Method)*/
b := bson.M{"$and": []bson.M{
	bson.M{"active": true},
	bson.M{"dynamic": false},
	bson.M{"origin":
	    bson.M{"$in":
		[]interface{}{origin, false},
	    },
	},
	bson.M{"host":
	    bson.M{"$in":
		[]interface{}{req.Host, false},
	    },
	},
	bson.M{"method":
	    bson.M{"$in":
		[]interface{}{req.Method, false},
	    },
	},
	bson.M{"path":
	    bson.M{"$in":
		[]interface{}{req.URL.Path, false},
	    },
	},
	bson.M{"query":
	    bson.M{"$in":
		[]interface{}{req.URL.Query().Encode(), false},
	    },
	},
    }}

			//b := bson.M{"active": true, "dynamic": false, "host": req.Host, "method": req.Method, "path": req.URL.Path, "query": req.URL.Query().Encode()}
			err := rules.Find(b).One(&rule)//.Sort("priority")
			//log.Printf("Query: %+v, Res: %+v", b, rule)
			if err == nil {
			    status, err := strconv.Atoi(rule.Status)
			    reqbody := ioutil.NopCloser(bytes.NewBufferString(rule.ReqBody))
			    respbody := ioutil.NopCloser(bytes.NewBufferString(rule.Body))
			    log.Printf("%+v", rule)
			    resp := NewResponse(req, rule.RespHeader, status, respbody)
			    ctx.UserData = ContextUserData{Store: true, Time: 0, Body: reqbody, Header: rule.RespHeader, Origin: origin}
			    return req, resp
			}

			// read the whole body
			reqbody, err = ioutil.ReadAll(req.Body)
			if err != nil {
				ctx.Warnf("Cannot read request body %s", err)
			}

			defer req.Body.Close()
			req.Body = ioutil.NopCloser(bytes.NewBuffer(reqbody))

			bodyreader = bytes.NewReader(reqbody)

		} else {
		    bodyreader = req.Body
		}

		ctx.UserData = ContextUserData{Store: true, Time: time.Now().UnixNano(), Body: bodyreader, Header: req.Header, Origin: origin}
		return req, nil
	})

	proxy.OnResponse().DoFunc(func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
		//ctx.Logf("Method: %s - host: %s", ctx.Resp.Request.Method, ctx.Resp.Request.Host)
		if c.Database != nil && ctx.UserData != nil && ctx.UserData.(ContextUserData).Store && ctx.Resp.Request.Method != "CONNECT" {
			// get response content type
			respctype := getContentType(ctx.Resp.Header.Get("Content-Type"))

			//log.Printf("Resp Contenttype %s", respctype)

			respid := bson.NewObjectId()
			//log.Printf("Resp id: %s, host: %s", respid.Hex(), ctx.Resp.Request.Host)

			filename := filepath.Join(tmpdir, respid.Hex())

			//log.Printf("Duplicating Body file id: %s", respid.String())
			resp.Body = NewTeeReadCloser(resp.Body, NewFileStream(filename, *db, respctype, respid, ctx))

			reqctype := getContentType(ctx.Resp.Request.Header.Get("Content-Type"))

			//log.Printf("Req Contenttype %s", reqctype)

			if reqctype == "application/x-www-form-urlencoded" {
				//log.Printf("setting req content type to text/plain for saving to mongo")
				reqctype = "text/plain"
			}

			reqid := bson.NewObjectId()
			//log.Printf("Req id: %s, host: %s", reqid.Hex(), ctx.Resp.Request.Host)

			saveFileToMongo(*db, reqid, reqctype, ctx.UserData.(ContextUserData).Body, reqid.Hex(), ctx)

			// prepare document
			content := Content{
				//Id: docid,
				Request: Request{
					Origin:	 ctx.UserData.(ContextUserData).Origin,
					Path:    ctx.Resp.Request.URL.Path,
					Query:   ctx.Resp.Request.URL.Query().Encode(),
					FileId:  reqid,
					Host:    ctx.Resp.Request.Host,
					Method:  ctx.Resp.Request.Method,
					Time:    float32(time.Now().UnixNano()-ctx.UserData.(ContextUserData).Time) / 1.0e9,
					Headers: ctx.UserData.(ContextUserData).Header},
				Response: Response{
					Status:  ctx.Resp.StatusCode,
					Headers: ctx.Resp.Header,
					FileId:  respid},
				SocketUUID: ctx.Uuid,
				Date: time.Now(),
			}

			err := c.Insert(content)
			if err != nil {
				ctx.Logf("Can't insert document: %v\n", err)
			}

		}
		return resp
	})

	log.Println("Starting Proxy")
	log.Fatalln(http.ListenAndServe(*addr, proxy))
}

type TeeReadCloser struct {
	r io.Reader
	w io.WriteCloser
	c io.Closer
}

func NewTeeReadCloser(r io.ReadCloser, w io.WriteCloser) io.ReadCloser {
	return &TeeReadCloser{io.TeeReader(r, w), w, r}
}

func (t *TeeReadCloser) Read(b []byte) (int, error) {
	return t.r.Read(b)
}

func (t *TeeReadCloser) Close() error {
	err1 := t.c.Close()
	err2 := t.w.Close()
	if err1 == nil && err2 == nil {
		return nil
	}
	if err1 != nil {
		return err2
	}
	return err1
}

type FileStream struct {
	path        string
	db          mgo.Database
	contentType string
	objectId    bson.ObjectId
	f           *os.File
	ctx	    *goproxy.ProxyCtx
}

func NewFileStream(path string, db mgo.Database, contentType string, objectId bson.ObjectId, ctx *goproxy.ProxyCtx) *FileStream {
    return &FileStream{path: path, db: db, contentType: contentType, objectId: objectId, f: nil, ctx: ctx}
}

func (fs *FileStream) Write(b []byte) (nr int, err error) {
	if fs.f == nil {
		fs.f, err = os.Create(fs.path)
		if err != nil {
			return 0, err
		}
	}
	return fs.f.Write(b)
}

func (fs *FileStream) Close() error {
	if fs.f == nil {
		return errors.New("FileStream was never written into")
	}
	fs.f.Seek(0, 0)
	saveFileToMongo(fs.db, fs.objectId, fs.contentType, fs.f, fs.objectId.Hex(), fs.ctx)
	err := fs.f.Close()
	if err == nil {
		err2 := os.Remove(fs.path)
		if err2 != nil {
			fs.ctx.Logf("Unable to delete file")
		}
	}
	return err
}

func getMongoFileContent(db mgo.Database, objId bson.ObjectId) (file *mgo.GridFile, err error) {
    file, err = db.GridFS("fs").OpenId(objId)

    if err != nil {
	return file, err
	if err == mgo.ErrNotFound {
	}
    }
    //defer file.Close()

    return file, err
}

// Store file in MongoDB GridFS
func saveFileToMongo(db mgo.Database, objId bson.ObjectId, contentType string, openFile io.Reader, fileName string, ctx *goproxy.ProxyCtx) {
	mdbfile, err := db.GridFS("fs").Create(fileName)
	if err == nil {
		mdbfile.SetContentType(contentType)
		mdbfile.SetId(objId)
		_, err = io.Copy(mdbfile, openFile)
		if err != nil {
			ctx.Logf("Unable to copy to mongo: %s - %v", fileName, err)
		}
		err = mdbfile.Close()
		if err != nil {
			ctx.Logf("Unable to close copy to mongo")
		}
		ctx.Logf("MongoDB body file saved")
	}
}

func getContentType(s string) string {
	arr := strings.Split(s, ";")
	return arr[0]
}
func ipAddrFromRemoteAddr(s string) string {
        idx := strings.LastIndex(s, ":")
        if idx == -1 {
                return s
        }
        return s[:idx]
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/gorilla/mux"
)

func index(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintln(w, "Welcome!")
}

func todoIndex(w http.ResponseWriter, r *http.Request) {
	todos := Todos{
		Todo{Name: "Write presentation"},
		Todo{Name: "Host meetup"},
	}

	if err := json.NewEncoder(w).Encode(todos); err != nil {
		panic(err)
	}
}

func todoShow(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	todoId := vars["todoId"]
	fmt.Fprintln(w, "Todo show:", todoId)
}
<commit_msg>Better instruct HTTP response:<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/gorilla/mux"
)

func index(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintln(w, "Welcome!")
}

func todoIndex(w http.ResponseWriter, r *http.Request) {
	todos := Todos{
		Todo{Name: "Write presentation"},
		Todo{Name: "Host meetup"},
	}

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(todos); err != nil {
		panic(err)
	}
}

func todoShow(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	todoId := vars["todoId"]
	fmt.Fprintln(w, "Todo show:", todoId)
}
<|endoftext|>


<commit_before>/* Plugin provides all the required interfaces to implement a GoPlugin
 *
 * Available API :
 *
 * func PluginInit(pluginImplConf PluginImplConf) (*Plugin, error)
 * >> Initialize a Plugin with specified Configuration
 *
 * func (plugin *Plugin) RegisterFunc(funcName string, method func([]byte) []byte)
 * >> Register a method to be executed for a Specified Path
 *
 * func (plugin *Plugin) Start() error
 * >> Start the execution of the specifiec Plugin
 *
 * func (plugin *Plugin) Stop() error
 * >> Stop the execution of that Plugin.
 */

package GoPlug

import (
	"com.ss/goplugin/PluginConn"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

/* The plugin impleentaion configuration  */
type PluginImplConf struct {
	// Plugin location path
	PluginLoc string
	// The Name of the plugin
	Name string
	// The namespace of the plugin [optional - default: nil]
	Namespace string
	// The URL to reach the plugin over http (ie unix://ExamplePlug) [optional - default: unix://<Namespace><Name>]
	Url string
	// The LazyLoad configuration [optional - default: false]
	LazyLoad bool
	// The Function that would be called on Plugin Activation
	Activator func([]byte) []byte
	// The Function that would be called on Plugin DeActivation
	Stopper func([]byte) []byte
}

type PluginImpl struct {
	pluginServer   *PluginConn.PluginServer
	methodRegistry map[string]func([]byte) []byte
	sockFile       string
	addr           string
}

/* Init a plugin for a specific Plugin Conf */
func PluginInit(pluginImplConf PluginImplConf) (*PluginImpl, error) {

	plugin := &PluginImpl{}
	pluginConf := PluginConf{}

	// Check pluginImplConf
	if pluginImplConf.PluginLoc == "" {
		return nil, fmt.Errorf("Invalid Configuration : PluginLoc file should be specified")
	}

	// Check name
	if pluginImplConf.Name == "" {
		return nil, fmt.Errorf("Invalid Configuration : Name should be specified")
	}
	pluginConf.Name = pluginImplConf.Name
	pluginConf.NameSpace = pluginImplConf.Namespace

	// Check url
	pluginConf.Url = pluginImplConf.Url
	if pluginImplConf.Url == "" {
		pluginConf.Url = "unix://" + pluginConf.NameSpace + pluginConf.Name
	}

	// Get conf file and Sock
	confFile := filepath.Join(pluginImplConf.PluginLoc, pluginConf.NameSpace+pluginConf.Name+".pconf")
	pwd, _ := os.Getwd()
	sockFileLoc := filepath.Join(pwd, pluginImplConf.PluginLoc)
	pluginConf.Sock = filepath.Join(sockFileLoc, pluginConf.NameSpace+pluginConf.Name+".sock")

	// Get Lazyload
	pluginConf.LazyLoad = pluginImplConf.LazyLoad

	// Load Plugin Configuration
	confSaveError := saveConfigs(confFile, pluginConf)
	if confSaveError != nil {
		fmt.Println("Configuration save failed to the file: ", confFile, ", Error: ", confSaveError)
		return nil, fmt.Errorf("Failed to save Configuration")
	}
	plugin.sockFile = pluginConf.Sock
	plugin.addr = pluginConf.Url

	// Initiate the Method Registry
	plugin.methodRegistry = make(map[string]func([]byte) []byte)

	plugin.methodRegistry["Activate"] = pluginImplConf.Activator
	plugin.methodRegistry["Stop"] = pluginImplConf.Stopper

	return plugin, nil
}

/* Internal Method: To Register method for the Plugin */
func (plugin *PluginImpl) Register() {

	http.Handle("/", plugin)
}

/* Internal Method: To handle all http request */
func (plugin *PluginImpl) ServeHTTP(res http.ResponseWriter, req *http.Request) {

	methodName := strings.Split(req.URL.Path, "/")[1]
	//fmt.Printf("URL found: %s\n", methodName)
	if methodName == "" {
		res.WriteHeader(400)
	} else {
		method, ok := plugin.methodRegistry[methodName]
		if ok {
			// Check if the method is Activate
			if methodName == "Activate" {
				methodReg := plugin.methodRegistry
				methods := make([]string, len(methodReg))
				idx := 0
				for key, _ := range methodReg {
					methods[idx] = key
					idx++
				}
				data, marshalErr := json.Marshal(methods)
				if marshalErr != nil {
					fmt.Println("failed to marshal methods")
					res.WriteHeader(400)
				}
				// Write the methods list
				res.Write(data)
			}

			defer req.Body.Close()
			body, _ := ioutil.ReadAll(req.Body)
			returnData := method(body)
			if returnData != nil {
				res.Write(returnData)
			}
			res.WriteHeader(200)
		} else {
			res.WriteHeader(400)
		}
	}
}

/* Method to register function for the plugin */
func (plugin *PluginImpl) RegisterFunc(funcName string, method func([]byte) []byte) {
	plugin.methodRegistry[funcName] = method
}

/* Start the Plugin Service */
func (plugin *PluginImpl) Start() error {

	sockFile := plugin.sockFile
	addr := plugin.addr
	// Create the Plugin Server
	config := &PluginConn.ServerConfiguration{Registrar: plugin, SockFile: sockFile, Addr: addr}
	server, err := PluginConn.NewPluginServer(config)
	if err != nil {
		fmt.Printf("Failed to Create server\n")
		return fmt.Errorf("Failed to Create the server")
	}
	plugin.pluginServer = server

	plugin.pluginServer.Start()

	return nil
}

/* Stop the Plugin service */
func (plugin *PluginImpl) Stop() error {
	err := plugin.pluginServer.Shutdown()
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Added auto deletion of the pconf file<commit_after>/* Plugin provides all the required interfaces to implement a GoPlugin
 *
 * Available API :
 *
 * func PluginInit(pluginImplConf PluginImplConf) (*Plugin, error)
 * >> Initialize a Plugin with specified Configuration
 *
 * func (plugin *Plugin) RegisterFunc(funcName string, method func([]byte) []byte)
 * >> Register a method to be executed for a Specified Path
 *
 * func (plugin *Plugin) Start() error
 * >> Start the execution of the specifiec Plugin
 *
 * func (plugin *Plugin) Stop() error
 * >> Stop the execution of that Plugin.
 */

package GoPlug

import (
	"com.ss/goplugin/PluginConn"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

/* The plugin impleentaion configuration  */
type PluginImplConf struct {
	// Plugin location path
	PluginLoc string
	// The Name of the plugin
	Name string
	// The namespace of the plugin [optional - default: nil]
	Namespace string
	// The URL to reach the plugin over http (ie unix://ExamplePlug) [optional - default: unix://<Namespace><Name>]
	Url string
	// The LazyLoad configuration [optional - default: false]
	LazyLoad bool
	// The Function that would be called on Plugin Activation
	Activator func([]byte) []byte
	// The Function that would be called on Plugin DeActivation
	Stopper func([]byte) []byte
}

type PluginImpl struct {
	pluginServer   *PluginConn.PluginServer
	methodRegistry map[string]func([]byte) []byte
	sockFile       string
	addr           string
	confFile       string
}

/* Init a plugin for a specific Plugin Conf */
func PluginInit(pluginImplConf PluginImplConf) (*PluginImpl, error) {

	plugin := &PluginImpl{}
	pluginConf := PluginConf{}

	// Check pluginImplConf
	if pluginImplConf.PluginLoc == "" {
		return nil, fmt.Errorf("Invalid Configuration : PluginLoc file should be specified")
	}

	// Check name
	if pluginImplConf.Name == "" {
		return nil, fmt.Errorf("Invalid Configuration : Name should be specified")
	}
	pluginConf.Name = pluginImplConf.Name
	pluginConf.NameSpace = pluginImplConf.Namespace

	// Check url
	pluginConf.Url = pluginImplConf.Url
	if pluginImplConf.Url == "" {
		pluginConf.Url = "unix://" + pluginConf.NameSpace + pluginConf.Name
	}

	// Get conf file and Sock
	confFile := filepath.Join(pluginImplConf.PluginLoc, pluginConf.NameSpace+pluginConf.Name+".pconf")
	pwd, _ := os.Getwd()
	sockFileLoc := filepath.Join(pwd, pluginImplConf.PluginLoc)
	pluginConf.Sock = filepath.Join(sockFileLoc, pluginConf.NameSpace+pluginConf.Name+".sock")

	// Get Lazyload
	pluginConf.LazyLoad = pluginImplConf.LazyLoad

	// Load Plugin Configuration
	confSaveError := saveConfigs(confFile, pluginConf)
	if confSaveError != nil {
		fmt.Println("Configuration save failed to the file: ", confFile, ", Error: ", confSaveError)
		return nil, fmt.Errorf("Failed to save Configuration")
	}
	plugin.sockFile = pluginConf.Sock
	plugin.addr = pluginConf.Url

	// Initiate the Method Registry
	plugin.methodRegistry = make(map[string]func([]byte) []byte)

	plugin.methodRegistry["Activate"] = pluginImplConf.Activator
	plugin.methodRegistry["Stop"] = pluginImplConf.Stopper

	plugin.confFile = confFile

	return plugin, nil
}

/* Internal Method: To Register method for the Plugin */
func (plugin *PluginImpl) Register() {

	http.Handle("/", plugin)
}

/* Internal Method: To handle all http request */
func (plugin *PluginImpl) ServeHTTP(res http.ResponseWriter, req *http.Request) {

	methodName := strings.Split(req.URL.Path, "/")[1]
	//fmt.Printf("URL found: %s\n", methodName)
	if methodName == "" {
		res.WriteHeader(400)
	} else {
		method, ok := plugin.methodRegistry[methodName]
		if ok {
			// Check if the method is Activate
			if methodName == "Activate" {
				methodReg := plugin.methodRegistry
				methods := make([]string, len(methodReg))
				idx := 0
				for key, _ := range methodReg {
					methods[idx] = key
					idx++
				}
				data, marshalErr := json.Marshal(methods)
				if marshalErr != nil {
					fmt.Println("failed to marshal methods")
					res.WriteHeader(400)
				}
				// Write the methods list
				res.Write(data)
			}

			defer req.Body.Close()
			body, _ := ioutil.ReadAll(req.Body)
			returnData := method(body)
			if returnData != nil {
				res.Write(returnData)
			}
			res.WriteHeader(200)
		} else {
			res.WriteHeader(400)
		}
	}
}

/* Method to register function for the plugin */
func (plugin *PluginImpl) RegisterFunc(funcName string, method func([]byte) []byte) {
	plugin.methodRegistry[funcName] = method
}

/* Start the Plugin Service */
func (plugin *PluginImpl) Start() error {

	sockFile := plugin.sockFile
	addr := plugin.addr
	// Create the Plugin Server
	config := &PluginConn.ServerConfiguration{Registrar: plugin, SockFile: sockFile, Addr: addr}
	server, err := PluginConn.NewPluginServer(config)
	if err != nil {
		fmt.Printf("Failed to Create server\n")
		return fmt.Errorf("Failed to Create the server")
	}
	plugin.pluginServer = server

	plugin.pluginServer.Start()

	return nil
}

/* Stop the Plugin service */
func (plugin *PluginImpl) Stop() error {
	err := plugin.pluginServer.Shutdown()
	if err != nil {
		return err
	}
	err = os.Remove(plugin.confFile)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/flynn/go-discoverd"
	"github.com/flynn/go-flynn/postgres"
)

var storageDir = flag.String("s", "", "Path to store files, instead of Postgres")

func errorResponse(w http.ResponseWriter, err error) {
	if err == ErrNotFound {
		http.Error(w, "NotFound", 404)
		return
	}
	log.Println("error:", err)
	http.Error(w, "Internal Server Error", 500)
}

type File interface {
	io.ReadSeeker
	io.Closer
	Size() int64
	ModTime() time.Time
	Type() string
	ETag() string
}

type Filesystem interface {
	Open(name string) (File, error)
	Put(name string, r io.Reader, typ string) error
	Delete(name string) error
}

var ErrNotFound = errors.New("file not found")

func handler(fs Filesystem) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		switch req.Method {
		case "HEAD", "GET":
			file, err := fs.Open(req.URL.Path)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			log.Println("GET", req.RequestURI)
			w.Header().Set("Content-Length", strconv.FormatInt(file.Size(), 10))
			w.Header().Set("Content-Type", file.Type())
			w.Header().Set("Etag", file.ETag())
			http.ServeContent(w, req, req.URL.Path, file.ModTime(), file)
		case "PUT":
			err := fs.Put(req.URL.Path, req.Body, req.Header.Get("Content-Type"))
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("PUT", req.RequestURI)
		case "DELETE":
			err := fs.Delete(req.URL.Path)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("DELETE", req.RequestURI)
		default:
			w.WriteHeader(http.StatusMethodNotAllowed)
		}
	})
}

func main() {
	flag.Parse()

	addr := os.Getenv("PORT")
	if addr == "" {
		addr = "3001"
	}
	addr = ":" + addr

	var fs Filesystem
	var storageDesc string

	if *storageDir != "" {
		fs = NewOSFilesystem(*storageDir)
		storageDesc = *storageDir
	} else {
		db, err := postgres.Open("", "")
		if err != nil {
			log.Fatal(err)
		}
		fs, err = NewPostgresFilesystem(db.DB)
		if err != nil {
			log.Fatal(err)
		}
		storageDesc = "Postgres"

		if err := discoverd.Register("shelf", addr); err != nil {
			log.Fatal(err)
		}
	}

	log.Println("Shelf serving files on " + addr + " from " + storageDesc)
	log.Fatal(http.ListenAndServe(addr, handler(fs)))
}
<commit_msg>Implement support for -p flag that defines server listen port<commit_after>package main

import (
	"errors"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/flynn/go-discoverd"
	"github.com/flynn/go-flynn/postgres"
)

var storageDir = flag.String("s", "", "Path to store files, instead of Postgres")
var listenPort = flag.String("p", "3001", "Port to listen on")

func errorResponse(w http.ResponseWriter, err error) {
	if err == ErrNotFound {
		http.Error(w, "NotFound", 404)
		return
	}
	log.Println("error:", err)
	http.Error(w, "Internal Server Error", 500)
}

type File interface {
	io.ReadSeeker
	io.Closer
	Size() int64
	ModTime() time.Time
	Type() string
	ETag() string
}

type Filesystem interface {
	Open(name string) (File, error)
	Put(name string, r io.Reader, typ string) error
	Delete(name string) error
}

var ErrNotFound = errors.New("file not found")

func handler(fs Filesystem) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		switch req.Method {
		case "HEAD", "GET":
			file, err := fs.Open(req.URL.Path)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			log.Println("GET", req.RequestURI)
			w.Header().Set("Content-Length", strconv.FormatInt(file.Size(), 10))
			w.Header().Set("Content-Type", file.Type())
			w.Header().Set("Etag", file.ETag())
			http.ServeContent(w, req, req.URL.Path, file.ModTime(), file)
		case "PUT":
			err := fs.Put(req.URL.Path, req.Body, req.Header.Get("Content-Type"))
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("PUT", req.RequestURI)
		case "DELETE":
			err := fs.Delete(req.URL.Path)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("DELETE", req.RequestURI)
		default:
			w.WriteHeader(http.StatusMethodNotAllowed)
		}
	})
}

func main() {
	flag.Parse()

	addr := os.Getenv("PORT")
	if addr == "" {
		addr = *listenPort
	}
	addr = ":" + addr

	var fs Filesystem
	var storageDesc string

	if *storageDir != "" {
		fs = NewOSFilesystem(*storageDir)
		storageDesc = *storageDir
	} else {
		db, err := postgres.Open("", "")
		if err != nil {
			log.Fatal(err)
		}
		fs, err = NewPostgresFilesystem(db.DB)
		if err != nil {
			log.Fatal(err)
		}
		storageDesc = "Postgres"

		if err := discoverd.Register("shelf", addr); err != nil {
			log.Fatal(err)
		}
	}

	log.Println("Shelf serving files on " + addr + " from " + storageDesc)
	log.Fatal(http.ListenAndServe(addr, handler(fs)))
}
<|endoftext|>