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

import (
	"encoding/json"
	"fmt"
	"hash/crc64"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/googlecodelabs/tools/claat/fetch/drive/auth"
	"github.com/googlecodelabs/tools/claat/parser"
	"github.com/googlecodelabs/tools/claat/types"
	"github.com/googlecodelabs/tools/claat/util"
)

const (
	// supported codelab source types must be registered parsers
	// TODO: define these in claat/parser/..., e.g. in parser/gdoc
	SrcInvalid   srcType = ""
	SrcGoogleDoc srcType = "gdoc" // Google Docs doc
	SrcMarkdown  srcType = "md"   // Markdown text

	// driveAPI is a base URL for Drive API
	driveAPI = "https://www.googleapis.com/drive/v3"
)

// srcType is codelab source type
type srcType string

// resource is a codelab resource, loaded from local file
// or fetched from remote location.
type resource struct {
	typ  srcType       // source type
	body io.ReadCloser // resource body
	mod  time.Time     // last update of content
}

// codelab wraps types.Codelab, while adding source type
// and modified timestamp fields.
type codelab struct {
	*types.Codelab
	Typ srcType   //  source type
	Mod time.Time // last modified timestamp
}

type Fetcher struct {
	authHelper   *auth.Helper
	authToken    string
	crcTable     *crc64.Table
	passMetadata map[string]bool
}

func NewFetcher(at string, pm map[string]bool) (*Fetcher, error) {
	h, err := auth.NewHelper(at, auth.ProviderGoogle, nil)
	if err != nil {
		return nil, err
	}

	return &Fetcher{
		authHelper:   h,
		authToken:    at,
		crcTable:     crc64.MakeTable(crc64.ECMA),
		passMetadata: pm,
	}, nil
}

// SlurpCodelab retrieves and parses codelab source.
// It takes the source, plus an auth token and a set of extra metadata to pass along.
// It returns parsed codelab and its source type.
//
// The function will also fetch and parse fragments included
// with types.ImportNode.
func (f *Fetcher) SlurpCodelab(src string) (*codelab, error) {
	res, err := f.fetch(src)
	if err != nil {
		return nil, err
	}
	defer res.body.Close()

	opts := *parser.NewOptions()
	opts.PassMetadata = f.passMetadata

	clab, err := parser.Parse(string(res.typ), res.body, opts)
	if err != nil {
		return nil, err
	}

	// fetch imports and parse them as fragments
	var imports []*types.ImportNode
	for _, st := range clab.Steps {
		imports = append(imports, types.ImportNodes(st.Content.Nodes)...)
	}
	ch := make(chan error, len(imports))
	defer close(ch)
	for _, imp := range imports {
		go func(n *types.ImportNode) {
			frag, err := f.slurpFragment(n.URL)
			if err != nil {
				ch <- fmt.Errorf("%s: %v", n.URL, err)
				return
			}
			n.Content.Nodes = frag
			ch <- nil
		}(imp)
	}
	for range imports {
		if err := <-ch; err != nil {
			return nil, err
		}
	}

	v := &codelab{
		Codelab: clab,
		Typ:     res.typ,
		Mod:     res.mod,
	}
	return v, nil
}

func (f *Fetcher) SlurpImages(src, dir string, steps []*types.Step) (map[string]string, error) {
	// make sure img dir exists
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}

	type res struct {
		url, file string
		err       error
	}

	ch := make(chan *res, 100)
	defer close(ch)
	var count int
	for _, st := range steps {
		nodes := types.ImageNodes(st.Content.Nodes)
		count += len(nodes)
		for _, n := range nodes {
			go func(n *types.ImageNode) {
				url := n.Src
				file, err := f.slurpBytes(src, dir, url)
				if err == nil {
					n.Src = filepath.Join(util.ImgDirname, file)
				}
				ch <- &res{url, file, err}
			}(n)
		}
	}

	imap := make(map[string]string, count)
	var err error
	for i := 0; i < count; i++ {
		r := <-ch
		imap[r.file] = r.url
		if r.err != nil && err == nil {
			// record first error
			err = fmt.Errorf("%s => %s: %v", r.url, r.file, r.err)
		}
	}

	return imap, err
}

func (f *Fetcher) slurpBytes(codelabSrc, dir, imgURL string) (string, error) {
	// images can be local in Markdown cases or remote.
	// Only proceed a simple copy on local reference.
	var b []byte
	var ext string
	u, err := url.Parse(imgURL)
	if err != nil {
		return "", err
	}

	// If the codelab source is being downloaded from the network, then we should interpret
	// the image URL in the same way.
	srcUrl, err := url.Parse(codelabSrc)
	if err == nil && srcUrl.Host != "" {
		u = srcUrl.ResolveReference(u)
	}

	if u.Host == "" {
		if imgURL, err = restrictPathToParent(imgURL, filepath.Dir(codelabSrc)); err != nil {
			return "", err
		}
		b, err = ioutil.ReadFile(imgURL)
		ext = filepath.Ext(imgURL)
	} else {
		b, err = slurpRemoteBytes(f.authHelper.DriveClient(), u.String(), 5)
		if string(b[6:10]) == "JFIF" {
			ext = ".jpeg"
		} else if string(b[0:3]) == "GIF" {
			ext = ".gif"
		} else {
			ext = ".png"
		}
	}
	if err != nil {
		return "", err
	}

	crc := crc64.Checksum(b, f.crcTable)
	file := fmt.Sprintf("%x%s", crc, ext)
	dst := filepath.Join(dir, file)
	return file, ioutil.WriteFile(dst, b, 0644)
}

func (f *Fetcher) slurpFragment(url string) ([]types.Node, error) {
	res, err := fetchRemote(url, f.authToken, true)
	if err != nil {
		return nil, err
	}
	defer res.body.Close()
	return parser.ParseFragment(string(res.typ), res.body)
}

// fetch retrieves codelab doc either from local disk
// or a remote location.
// The caller is responsible for closing returned stream.
func (f *Fetcher) fetch(name string) (*resource, error) {
	fi, err := os.Stat(name)
	if os.IsNotExist(err) {
		return fetchRemote(name, f.authToken, false)
	}
	r, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	return &resource{
		body: r,
		typ:  SrcMarkdown,
		mod:  fi.ModTime(),
	}, nil
}

// fetchRemote retrieves resource r from the network.
//
// If urlStr is not a URL, i.e. does not have the host part, it is considered to be
// a Google Doc ID and fetched accordingly. Otherwise, a simple GET request
// is used to retrieve the contents.
//
// The caller is responsible for closing returned stream.
// If nometa is true, resource.mod may have zero value.
func fetchRemote(urlStr, authToken string, nometa bool) (*resource, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}
	if u.Host == "" || u.Host == "docs.google.com" {
		return fetchDriveFile(urlStr, authToken, nometa)
	}
	return fetchRemoteFile(urlStr)
}

// fetchRemoteFile retrieves codelab resource from url.
// It is a special case of fetchRemote function.
func fetchRemoteFile(url string) (*resource, error) {
	res, err := retryGet(nil, url, 3)
	if err != nil {
		return nil, err
	}
	t, err := http.ParseTime(res.Header.Get("last-modified"))
	if err != nil {
		t = time.Now()
	}
	return &resource{
		body: res.Body,
		mod:  t,
		typ:  SrcMarkdown,
	}, nil
}

// fetchDriveFile uses Drive API to retrieve HTML representation of a Google Doc.
// See https://developers.google.com/drive/web/manage-downloads#downloading_google_documents
// for more details.
//
// If nometa is true, resource.mod will have zero value.
func fetchDriveFile(id, authToken string, nometa bool) (*resource, error) {
	id = gdocID(id)
	exportURL := gdocExportURL(id)
	h, err := auth.NewHelper(authToken, auth.ProviderGoogle, nil)
	if err != nil {
		return nil, err
	}
	client := h.DriveClient()

	if nometa {
		res, err := retryGet(client, exportURL, 7)
		if err != nil {
			return nil, err
		}
		return &resource{body: res.Body, typ: SrcGoogleDoc}, nil
	}

	q := url.Values{
		"fields":             {"id,mimeType,modifiedTime"},
		"supportsTeamDrives": {"true"},
	}
	u := fmt.Sprintf("%s/files/%s?%s", driveAPI, id, q.Encode())
	res, err := retryGet(client, u, 7)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	meta := &struct {
		ID       string    `json:"id"`
		MimeType string    `json:"mimeType"`
		Modified time.Time `json:"modifiedTime"`
	}{}
	if err := json.NewDecoder(res.Body).Decode(meta); err != nil {
		return nil, err
	}
	if meta.MimeType != "application/vnd.google-apps.document" {
		return nil, fmt.Errorf("%s: invalid mime type: %s", id, meta.MimeType)
	}

	if res, err = retryGet(client, exportURL, 7); err != nil {
		return nil, err
	}
	return &resource{
		body: res.Body,
		mod:  meta.Modified,
		typ:  SrcGoogleDoc,
	}, nil
}

func slurpRemoteBytes(client *http.Client, url string, n int) ([]byte, error) {
	res, err := retryGet(client, url, n)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	return ioutil.ReadAll(res.Body)
}

// retryGet tries to GET specified url up to n times.
// Default client will be used if not provided.
func retryGet(client *http.Client, url string, n int) (*http.Response, error) {
	if client == nil {
		client = http.DefaultClient
	}
	for i := 0; i <= n; i++ {
		if i > 0 {
			t := time.Duration((math.Pow(2, float64(i)) + rand.Float64()) * float64(time.Second))
			time.Sleep(t)
		}
		res, err := client.Get(url)
		// return early with a good response
		// the rest is error handling
		if err == nil && res.StatusCode == http.StatusOK {
			return res, nil
		}

		// sometimes Drive API wouldn't even start a response,
		// we get net/http: TLS handshake timeout instead:
		// consider this a temporary failure and retry again
		if err != nil {
			continue
		}
		// otherwise, decode error response and check for "rate limit"
		defer res.Body.Close()
		var erres struct {
			Error struct {
				Errors []struct{ Reason string }
			}
		}
		b, _ := ioutil.ReadAll(res.Body)
		json.Unmarshal(b, &erres)
		var rateLimit bool
		for _, e := range erres.Error.Errors {
			if e.Reason == "rateLimitExceeded" || e.Reason == "userRateLimitExceeded" {
				rateLimit = true
				break
			}
		}
		// this is neither a rate limit error, nor a server error:
		// retrying is useless
		if !rateLimit && res.StatusCode < http.StatusInternalServerError {
			return nil, fmt.Errorf("fetch %s: %s; %s", url, res.Status, b)
		}
	}
	return nil, fmt.Errorf("%s: failed after %d retries", url, n)
}

func gdocID(url string) string {
	const s = "/document/d/"
	if i := strings.Index(url, s); i >= 0 {
		url = url[i+len(s):]
	}
	if i := strings.IndexRune(url, '/'); i > 0 {
		url = url[:i]
	}
	return url
}

func gdocExportURL(id string) string {
	return fmt.Sprintf("%s/files/%s/export?mimeType=text/html", driveAPI, id)
}

// restrictPathToParent will ensure that assetPath is in parent.
// It will thus return an absolute path to the asset.
func restrictPathToParent(assetPath, parent string) (string, error) {
	parent, err := filepath.Abs(parent)
	if err != nil {
		return "", err
	}
	if !filepath.IsAbs(assetPath) {
		assetPath = filepath.Join(parent, assetPath)
	}
	if !strings.HasPrefix(assetPath, parent) {
		return "", fmt.Errorf("%s isn't a subdirectory of %s", assetPath, parent)
	}
	return assetPath, nil
}
<commit_msg>Make fetchRemote a method of Fetcher.<commit_after>// Copyright 2016-2019 Google LLC. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package fetch

import (
	"encoding/json"
	"fmt"
	"hash/crc64"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/googlecodelabs/tools/claat/fetch/drive/auth"
	"github.com/googlecodelabs/tools/claat/parser"
	"github.com/googlecodelabs/tools/claat/types"
	"github.com/googlecodelabs/tools/claat/util"
)

const (
	// supported codelab source types must be registered parsers
	// TODO: define these in claat/parser/..., e.g. in parser/gdoc
	SrcInvalid   srcType = ""
	SrcGoogleDoc srcType = "gdoc" // Google Docs doc
	SrcMarkdown  srcType = "md"   // Markdown text

	// driveAPI is a base URL for Drive API
	driveAPI = "https://www.googleapis.com/drive/v3"
)

// srcType is codelab source type
type srcType string

// resource is a codelab resource, loaded from local file
// or fetched from remote location.
type resource struct {
	typ  srcType       // source type
	body io.ReadCloser // resource body
	mod  time.Time     // last update of content
}

// codelab wraps types.Codelab, while adding source type
// and modified timestamp fields.
type codelab struct {
	*types.Codelab
	Typ srcType   //  source type
	Mod time.Time // last modified timestamp
}

type Fetcher struct {
	authHelper   *auth.Helper
	authToken    string
	crcTable     *crc64.Table
	passMetadata map[string]bool
}

func NewFetcher(at string, pm map[string]bool) (*Fetcher, error) {
	h, err := auth.NewHelper(at, auth.ProviderGoogle, nil)
	if err != nil {
		return nil, err
	}

	return &Fetcher{
		authHelper:   h,
		authToken:    at,
		crcTable:     crc64.MakeTable(crc64.ECMA),
		passMetadata: pm,
	}, nil
}

// SlurpCodelab retrieves and parses codelab source.
// It takes the source, plus an auth token and a set of extra metadata to pass along.
// It returns parsed codelab and its source type.
//
// The function will also fetch and parse fragments included
// with types.ImportNode.
func (f *Fetcher) SlurpCodelab(src string) (*codelab, error) {
	res, err := f.fetch(src)
	if err != nil {
		return nil, err
	}
	defer res.body.Close()

	opts := *parser.NewOptions()
	opts.PassMetadata = f.passMetadata

	clab, err := parser.Parse(string(res.typ), res.body, opts)
	if err != nil {
		return nil, err
	}

	// fetch imports and parse them as fragments
	var imports []*types.ImportNode
	for _, st := range clab.Steps {
		imports = append(imports, types.ImportNodes(st.Content.Nodes)...)
	}
	ch := make(chan error, len(imports))
	defer close(ch)
	for _, imp := range imports {
		go func(n *types.ImportNode) {
			frag, err := f.slurpFragment(n.URL)
			if err != nil {
				ch <- fmt.Errorf("%s: %v", n.URL, err)
				return
			}
			n.Content.Nodes = frag
			ch <- nil
		}(imp)
	}
	for range imports {
		if err := <-ch; err != nil {
			return nil, err
		}
	}

	v := &codelab{
		Codelab: clab,
		Typ:     res.typ,
		Mod:     res.mod,
	}
	return v, nil
}

func (f *Fetcher) SlurpImages(src, dir string, steps []*types.Step) (map[string]string, error) {
	// make sure img dir exists
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}

	type res struct {
		url, file string
		err       error
	}

	ch := make(chan *res, 100)
	defer close(ch)
	var count int
	for _, st := range steps {
		nodes := types.ImageNodes(st.Content.Nodes)
		count += len(nodes)
		for _, n := range nodes {
			go func(n *types.ImageNode) {
				url := n.Src
				file, err := f.slurpBytes(src, dir, url)
				if err == nil {
					n.Src = filepath.Join(util.ImgDirname, file)
				}
				ch <- &res{url, file, err}
			}(n)
		}
	}

	imap := make(map[string]string, count)
	var err error
	for i := 0; i < count; i++ {
		r := <-ch
		imap[r.file] = r.url
		if r.err != nil && err == nil {
			// record first error
			err = fmt.Errorf("%s => %s: %v", r.url, r.file, r.err)
		}
	}

	return imap, err
}

func (f *Fetcher) slurpBytes(codelabSrc, dir, imgURL string) (string, error) {
	// images can be local in Markdown cases or remote.
	// Only proceed a simple copy on local reference.
	var b []byte
	var ext string
	u, err := url.Parse(imgURL)
	if err != nil {
		return "", err
	}

	// If the codelab source is being downloaded from the network, then we should interpret
	// the image URL in the same way.
	srcUrl, err := url.Parse(codelabSrc)
	if err == nil && srcUrl.Host != "" {
		u = srcUrl.ResolveReference(u)
	}

	if u.Host == "" {
		if imgURL, err = restrictPathToParent(imgURL, filepath.Dir(codelabSrc)); err != nil {
			return "", err
		}
		b, err = ioutil.ReadFile(imgURL)
		ext = filepath.Ext(imgURL)
	} else {
		b, err = slurpRemoteBytes(f.authHelper.DriveClient(), u.String(), 5)
		if string(b[6:10]) == "JFIF" {
			ext = ".jpeg"
		} else if string(b[0:3]) == "GIF" {
			ext = ".gif"
		} else {
			ext = ".png"
		}
	}
	if err != nil {
		return "", err
	}

	crc := crc64.Checksum(b, f.crcTable)
	file := fmt.Sprintf("%x%s", crc, ext)
	dst := filepath.Join(dir, file)
	return file, ioutil.WriteFile(dst, b, 0644)
}

func (f *Fetcher) slurpFragment(url string) ([]types.Node, error) {
	res, err := f.fetchRemote(url, true)
	if err != nil {
		return nil, err
	}
	defer res.body.Close()
	return parser.ParseFragment(string(res.typ), res.body)
}

// fetch retrieves codelab doc either from local disk
// or a remote location.
// The caller is responsible for closing returned stream.
func (f *Fetcher) fetch(name string) (*resource, error) {
	fi, err := os.Stat(name)
	if os.IsNotExist(err) {
		return f.fetchRemote(name, false)
	}
	r, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	return &resource{
		body: r,
		typ:  SrcMarkdown,
		mod:  fi.ModTime(),
	}, nil
}

// fetchRemote retrieves resource r from the network.
//
// If urlStr is not a URL, i.e. does not have the host part, it is considered to be
// a Google Doc ID and fetched accordingly. Otherwise, a simple GET request
// is used to retrieve the contents.
//
// The caller is responsible for closing returned stream.
// If nometa is true, resource.mod may have zero value.
func (f *Fetcher) fetchRemote(urlStr string, nometa bool) (*resource, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}
	if u.Host == "" || u.Host == "docs.google.com" {
		return fetchDriveFile(urlStr, f.authToken, nometa)
	}
	return fetchRemoteFile(urlStr)
}

// fetchRemoteFile retrieves codelab resource from url.
// It is a special case of fetchRemote function.
func fetchRemoteFile(url string) (*resource, error) {
	res, err := retryGet(nil, url, 3)
	if err != nil {
		return nil, err
	}
	t, err := http.ParseTime(res.Header.Get("last-modified"))
	if err != nil {
		t = time.Now()
	}
	return &resource{
		body: res.Body,
		mod:  t,
		typ:  SrcMarkdown,
	}, nil
}

// fetchDriveFile uses Drive API to retrieve HTML representation of a Google Doc.
// See https://developers.google.com/drive/web/manage-downloads#downloading_google_documents
// for more details.
//
// If nometa is true, resource.mod will have zero value.
func fetchDriveFile(id, authToken string, nometa bool) (*resource, error) {
	id = gdocID(id)
	exportURL := gdocExportURL(id)
	h, err := auth.NewHelper(authToken, auth.ProviderGoogle, nil)
	if err != nil {
		return nil, err
	}
	client := h.DriveClient()

	if nometa {
		res, err := retryGet(client, exportURL, 7)
		if err != nil {
			return nil, err
		}
		return &resource{body: res.Body, typ: SrcGoogleDoc}, nil
	}

	q := url.Values{
		"fields":             {"id,mimeType,modifiedTime"},
		"supportsTeamDrives": {"true"},
	}
	u := fmt.Sprintf("%s/files/%s?%s", driveAPI, id, q.Encode())
	res, err := retryGet(client, u, 7)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	meta := &struct {
		ID       string    `json:"id"`
		MimeType string    `json:"mimeType"`
		Modified time.Time `json:"modifiedTime"`
	}{}
	if err := json.NewDecoder(res.Body).Decode(meta); err != nil {
		return nil, err
	}
	if meta.MimeType != "application/vnd.google-apps.document" {
		return nil, fmt.Errorf("%s: invalid mime type: %s", id, meta.MimeType)
	}

	if res, err = retryGet(client, exportURL, 7); err != nil {
		return nil, err
	}
	return &resource{
		body: res.Body,
		mod:  meta.Modified,
		typ:  SrcGoogleDoc,
	}, nil
}

func slurpRemoteBytes(client *http.Client, url string, n int) ([]byte, error) {
	res, err := retryGet(client, url, n)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	return ioutil.ReadAll(res.Body)
}

// retryGet tries to GET specified url up to n times.
// Default client will be used if not provided.
func retryGet(client *http.Client, url string, n int) (*http.Response, error) {
	if client == nil {
		client = http.DefaultClient
	}
	for i := 0; i <= n; i++ {
		if i > 0 {
			t := time.Duration((math.Pow(2, float64(i)) + rand.Float64()) * float64(time.Second))
			time.Sleep(t)
		}
		res, err := client.Get(url)
		// return early with a good response
		// the rest is error handling
		if err == nil && res.StatusCode == http.StatusOK {
			return res, nil
		}

		// sometimes Drive API wouldn't even start a response,
		// we get net/http: TLS handshake timeout instead:
		// consider this a temporary failure and retry again
		if err != nil {
			continue
		}
		// otherwise, decode error response and check for "rate limit"
		defer res.Body.Close()
		var erres struct {
			Error struct {
				Errors []struct{ Reason string }
			}
		}
		b, _ := ioutil.ReadAll(res.Body)
		json.Unmarshal(b, &erres)
		var rateLimit bool
		for _, e := range erres.Error.Errors {
			if e.Reason == "rateLimitExceeded" || e.Reason == "userRateLimitExceeded" {
				rateLimit = true
				break
			}
		}
		// this is neither a rate limit error, nor a server error:
		// retrying is useless
		if !rateLimit && res.StatusCode < http.StatusInternalServerError {
			return nil, fmt.Errorf("fetch %s: %s; %s", url, res.Status, b)
		}
	}
	return nil, fmt.Errorf("%s: failed after %d retries", url, n)
}

func gdocID(url string) string {
	const s = "/document/d/"
	if i := strings.Index(url, s); i >= 0 {
		url = url[i+len(s):]
	}
	if i := strings.IndexRune(url, '/'); i > 0 {
		url = url[:i]
	}
	return url
}

func gdocExportURL(id string) string {
	return fmt.Sprintf("%s/files/%s/export?mimeType=text/html", driveAPI, id)
}

// restrictPathToParent will ensure that assetPath is in parent.
// It will thus return an absolute path to the asset.
func restrictPathToParent(assetPath, parent string) (string, error) {
	parent, err := filepath.Abs(parent)
	if err != nil {
		return "", err
	}
	if !filepath.IsAbs(assetPath) {
		assetPath = filepath.Join(parent, assetPath)
	}
	if !strings.HasPrefix(assetPath, parent) {
		return "", fmt.Errorf("%s isn't a subdirectory of %s", assetPath, parent)
	}
	return assetPath, nil
}
<|endoftext|>


<commit_before>package calcium

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/docker/docker/pkg/archive"
	enginetypes "github.com/docker/engine-api/types"
	"gitlab.ricebook.net/platform/core/git"
	"gitlab.ricebook.net/platform/core/types"
	"gitlab.ricebook.net/platform/core/utils"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
)

// Get a random node from pod `podname`
func getRandomNode(c *Calcium, podname string) (*types.Node, error) {
	nodes, err := c.ListPodNodes(podname)
	if err != nil {
		return nil, err
	}

	nodemap := make(map[string]types.CPUMap)
	for _, n := range nodes {
		nodemap[n.Name] = n.CPU
	}
	nodename, err := c.scheduler.RandomNode(nodemap)
	if err != nil {
		return nil, err
	}

	return c.GetNode(podname, nodename)
}

// build image for repository
// since we wanna set UID for the user inside container, we have to know the uid parameter
//
// build directory is like:
//
//    buildDir ├─ :appname ├─ code
//             ├─ Dockerfile
//             ├─ launcher
//             ├─ launcheroot
func (c *Calcium) BuildImage(repository, version, uid string) (chan *types.BuildImageMessage, error) {
	ch := make(chan *types.BuildImageMessage)

	// use pod `dev` to build image
	node, err := getRandomNode(c, "dev")
	if err != nil {
		return ch, err
	}

	buildDir, err := ioutil.TempDir(os.TempDir(), "corebuild-")
	if err != nil {
		return ch, err
	}
	defer os.RemoveAll(buildDir)

	// parse repository name
	// code locates under /:repositoryname
	reponame, err := utils.GetGitRepoName(repository)
	if err != nil {
		return ch, err
	}

	// clone code into cloneDir
	// which is under buildDir and named as repository name
	cloneDir := filepath.Join(buildDir, reponame)
	if err := git.CloneRepository(repository, cloneDir, version, c.config.Git.PublicKey, c.config.Git.PrivateKey); err != nil {
		return ch, err
	}

	// ensure .git directory is removed
	// we don't want any history files to be retrieved
	if err := os.RemoveAll(filepath.Join(cloneDir, ".git")); err != nil {
		return ch, err
	}

	// use app.yaml file to create Specs instance
	// which we'll need to create Dockerfile later
	bytes, err := ioutil.ReadFile(filepath.Join(cloneDir, "app.yaml"))
	if err != nil {
		return ch, err
	}
	specs := types.Specs{}
	if err := yaml.Unmarshal(bytes, &specs); err != nil {
		return ch, err
	}

	// create launcher scripts and dockerfile
	if err := createLauncher(buildDir, specs); err != nil {
		return ch, err
	}
	if err := createDockerfile(buildDir, uid, reponame, specs); err != nil {
		return ch, err
	}

	// tag of image, later this will be used to push image to hub
	tag := fmt.Sprintf("%s/%s:%s", c.config.Docker.Hub, specs.Appname, utils.TruncateID(version))

	// create tar stream for Build API
	buildContext, err := createTarStream(buildDir)
	if err != nil {
		return ch, err
	}

	// must be put here because of that `defer os.RemoveAll(buildDir)`
	buildOptions := enginetypes.ImageBuildOptions{
		Tags:           []string{tag},
		SuppressOutput: false,
		NoCache:        true,
		Remove:         true,
		ForceRemove:    true,
		PullParent:     true,
	}
	resp, err := node.Engine.ImageBuild(context.Background(), buildContext, buildOptions)
	if err != nil {
		return ch, err
	}

	go func() {
		defer resp.Body.Close()
		decoder := json.NewDecoder(resp.Body)
		for {
			message := &types.BuildImageMessage{}
			err := decoder.Decode(message)
			if err != nil {
				if err == io.EOF {
					break
				} else {
					close(ch)
					return
				}
			}
			ch <- message
		}

		// About this "Khadgar", https://github.com/docker/docker/issues/10983#issuecomment-85892396
		// Just because Ben Schnetzer's cute Khadgar...
		rc, err := node.Engine.ImagePush(context.Background(), tag, enginetypes.ImagePushOptions{RegistryAuth: "Khadgar"})
		if err != nil {
			ch <- makeErrorBuildImageMessage(err)
			close(ch)
			return
		}

		defer rc.Close()
		decoder2 := json.NewDecoder(rc)
		for {
			message := &types.BuildImageMessage{}
			err := decoder2.Decode(message)
			if err != nil {
				close(ch)
				break
			}
			ch <- message
		}
	}()

	return ch, nil
}

func makeErrorBuildImageMessage(err error) *types.BuildImageMessage {
	return &types.BuildImageMessage{Error: err.Error()}
}

func createTarStream(path string) (io.ReadCloser, error) {
	tarOpts := &archive.TarOptions{
		ExcludePatterns: []string{},
		IncludeFiles:    []string{"."},
		Compression:     archive.Uncompressed,
		NoLchown:        true,
	}
	return archive.TarWithOptions(path, tarOpts)
}

// launcher scripts
// TODO use golang template
func createLauncher(buildDir string, specs types.Specs) error {
	entry := fmt.Sprintf("exec sudo -E -u %s $@", specs.Appname)
	entryRoot := "exec $@"

	tmpl := `#! /bin/sh
echo 32768 > /writable-proc/sys/net/core/somaxconn
echo 1 > /writable-proc/sys/vm/overcommit_memory
chmod 777 /dev/stdout
chmod 777 /dev/stderr

neednetwork=$1
if [ $neednetwork = "network" ]; then
    # wait for macvlan
    while ( ! ip addr show | grep 'UP' | grep 'vnbe'); do
        echo -n o
        sleep .5
    done
fi

sleep 1

shift

`
	f, err := os.Create(filepath.Join(buildDir, "launcher"))
	if err != nil {
		return err
	}
	defer f.Close()
	f.WriteString(tmpl)
	f.WriteString(entry)
	if err := f.Sync(); err != nil {
		return err
	}
	if err := f.Chmod(0755); err != nil {
		return err
	}

	fr, err := os.Create(filepath.Join(buildDir, "launcheroot"))
	if err != nil {
		return err
	}
	defer fr.Close()
	fr.WriteString(tmpl)
	fr.WriteString(entryRoot)
	if err := fr.Sync(); err != nil {
		return err
	}
	if err := fr.Chmod(0755); err != nil {
		return err
	}

	return nil
}

// Dockerfile
func createDockerfile(buildDir, uid, reponame string, specs types.Specs) error {
	f, err := os.Create(filepath.Join(buildDir, "Dockerfile"))
	if err != nil {
		return err
	}
	defer f.Close()

	f.WriteString(fmt.Sprintf("FROM %s\n", specs.Base))
	f.WriteString("ENV ERU 1\n")
	f.WriteString(fmt.Sprintf("ADD %s /%s\n", reponame, specs.Appname))
	f.WriteString("ADD launcher /usr/local/bin/launcher\n")
	f.WriteString("ADD launcheroot /usr/local/bin/launcheroot\n")
	f.WriteString(fmt.Sprintf("WORKDIR /%s\n", specs.Appname))
	f.WriteString(fmt.Sprintf("RUN useradd -u %s -d /nonexistent -s /sbin/nologin -U %s\n", uid, specs.Appname))
	for _, cmd := range specs.Build {
		f.WriteString(fmt.Sprintf("RUN %s\n", cmd))
	}

	if err := f.Sync(); err != nil {
		return err
	}
	return nil
}
<commit_msg>rewrite launcher using template<commit_after>package calcium

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"text/template"

	"github.com/docker/docker/pkg/archive"
	enginetypes "github.com/docker/engine-api/types"
	"gitlab.ricebook.net/platform/core/git"
	"gitlab.ricebook.net/platform/core/types"
	"gitlab.ricebook.net/platform/core/utils"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
)

const launcherScript = `#! /bin/sh
echo 32768 > /writable-proc/sys/net/core/somaxconn
echo 1 > /writable-proc/sys/vm/overcommit_memory
chmod 777 /dev/stdout
chmod 777 /dev/stderr

neednetwork=$1
if [ $neednetwork = "network" ]; then
    # wait for macvlan
    while ( ! ip addr show | grep 'UP' | grep 'vnbe'); do
        echo -n o
        sleep .5
    done
fi

sleep 1

shift

{{.Command}}
`

const dockerFile = `FROM {{.Base}}
ENV ERU 1
ADD %s /{{.Appname}}
ADD launcher /usr/local/bin/launcher
ADD launcheroot /usr/local/bin/launcheroot
WORKDIR /{{.Appname}}
RUN useradd -u %s -d /nonexistent -s /sbin/nologin -U {{.Appname}}
{{with .Build}}
{{range $index, $value := .}}
RUN {{$value}}
{{end}}
{{end}}
`

// Entry is used to format templates
type entry struct {
	Command string
}

// Get a random node from pod `podname`
func getRandomNode(c *Calcium, podname string) (*types.Node, error) {
	nodes, err := c.ListPodNodes(podname)
	if err != nil {
		return nil, err
	}

	nodemap := make(map[string]types.CPUMap)
	for _, n := range nodes {
		nodemap[n.Name] = n.CPU
	}
	nodename, err := c.scheduler.RandomNode(nodemap)
	if err != nil {
		return nil, err
	}

	return c.GetNode(podname, nodename)
}

// BuildImage will build image for repository
// since we wanna set UID for the user inside container, we have to know the uid parameter
//
// build directory is like:
//
//    buildDir ├─ :appname ├─ code
//             ├─ Dockerfile
//             ├─ launcher
//             ├─ launcheroot
func (c *Calcium) BuildImage(repository, version, uid string) (chan *types.BuildImageMessage, error) {
	ch := make(chan *types.BuildImageMessage)

	// use pod `dev` to build image
	node, err := getRandomNode(c, "dev")
	if err != nil {
		return ch, err
	}

	buildDir, err := ioutil.TempDir(os.TempDir(), "corebuild-")
	if err != nil {
		return ch, err
	}
	defer os.RemoveAll(buildDir)

	// parse repository name
	// code locates under /:repositoryname
	reponame, err := utils.GetGitRepoName(repository)
	if err != nil {
		return ch, err
	}

	// clone code into cloneDir
	// which is under buildDir and named as repository name
	cloneDir := filepath.Join(buildDir, reponame)
	if err := git.CloneRepository(repository, cloneDir, version, c.config.Git.PublicKey, c.config.Git.PrivateKey); err != nil {
		return ch, err
	}

	// ensure .git directory is removed
	// we don't want any history files to be retrieved
	if err := os.RemoveAll(filepath.Join(cloneDir, ".git")); err != nil {
		return ch, err
	}

	// use app.yaml file to create Specs instance
	// which we'll need to create Dockerfile later
	bytes, err := ioutil.ReadFile(filepath.Join(cloneDir, "app.yaml"))
	if err != nil {
		return ch, err
	}
	specs := types.Specs{}
	if err := yaml.Unmarshal(bytes, &specs); err != nil {
		return ch, err
	}

	// create launcher scripts and dockerfile
	if err := createLauncher(buildDir, specs); err != nil {
		return ch, err
	}
	if err := createDockerfile(buildDir, uid, reponame, specs); err != nil {
		return ch, err
	}

	// tag of image, later this will be used to push image to hub
	tag := fmt.Sprintf("%s/%s:%s", c.config.Docker.Hub, specs.Appname, utils.TruncateID(version))

	// create tar stream for Build API
	buildContext, err := createTarStream(buildDir)
	if err != nil {
		return ch, err
	}

	// must be put here because of that `defer os.RemoveAll(buildDir)`
	buildOptions := enginetypes.ImageBuildOptions{
		Tags:           []string{tag},
		SuppressOutput: false,
		NoCache:        true,
		Remove:         true,
		ForceRemove:    true,
		PullParent:     true,
	}
	resp, err := node.Engine.ImageBuild(context.Background(), buildContext, buildOptions)
	if err != nil {
		return ch, err
	}

	go func() {
		defer resp.Body.Close()
		decoder := json.NewDecoder(resp.Body)
		for {
			message := &types.BuildImageMessage{}
			err := decoder.Decode(message)
			if err != nil {
				if err == io.EOF {
					break
				} else {
					close(ch)
					return
				}
			}
			ch <- message
		}

		// About this "Khadgar", https://github.com/docker/docker/issues/10983#issuecomment-85892396
		// Just because Ben Schnetzer's cute Khadgar...
		rc, err := node.Engine.ImagePush(context.Background(), tag, enginetypes.ImagePushOptions{RegistryAuth: "Khadgar"})
		if err != nil {
			ch <- makeErrorBuildImageMessage(err)
			close(ch)
			return
		}

		defer rc.Close()
		decoder2 := json.NewDecoder(rc)
		for {
			message := &types.BuildImageMessage{}
			err := decoder2.Decode(message)
			if err != nil {
				close(ch)
				break
			}
			ch <- message
		}
	}()

	return ch, nil
}

func makeErrorBuildImageMessage(err error) *types.BuildImageMessage {
	return &types.BuildImageMessage{Error: err.Error()}
}

func createTarStream(path string) (io.ReadCloser, error) {
	tarOpts := &archive.TarOptions{
		ExcludePatterns: []string{},
		IncludeFiles:    []string{"."},
		Compression:     archive.Uncompressed,
		NoLchown:        true,
	}
	return archive.TarWithOptions(path, tarOpts)
}

// launcher scripts
func createLauncher(buildDir string, specs types.Specs) error {
	launcherScriptTemplate, _ := template.New("launcher script").Parse(launcherScript)

	entryCommand := fmt.Sprintf("exec sudo -E -u %s $@", specs.Appname)
	entryRootCommand := "exec $@"

	f, err := os.Create(filepath.Join(buildDir, "launcher"))
	if err != nil {
		return err
	}
	defer f.Close()
	launcherScriptTemplate.Execute(f, entry{Command: entryCommand})
	if err := f.Sync(); err != nil {
		return err
	}
	if err := f.Chmod(0755); err != nil {
		return err
	}

	fr, err := os.Create(filepath.Join(buildDir, "launcheroot"))
	if err != nil {
		return err
	}
	defer fr.Close()
	launcherScriptTemplate.Execute(fr, entry{Command: entryRootCommand})
	if err := fr.Sync(); err != nil {
		return err
	}
	if err := fr.Chmod(0755); err != nil {
		return err
	}

	return nil
}

// Dockerfile
func createDockerfile(buildDir, uid, reponame string, specs types.Specs) error {
	f, err := os.Create(filepath.Join(buildDir, "Dockerfile"))
	if err != nil {
		return err
	}
	defer f.Close()

	dockerFileFormatted := fmt.Sprintf(dockerFile, reponame, uid)
	t := template.New("docker file template")
	parsedTemplate, err := t.Parse(dockerFileFormatted)
	if err != nil {
		return err
	}
	parsedTemplate.Execute(f, specs)

	if err := f.Sync(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


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

// +build unit

package cmd

import (
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"testing"

	"github.com/control-center/serviced/cli/api"
	"github.com/control-center/serviced/domain/host"
	"github.com/control-center/serviced/domain/pool"
	"github.com/control-center/serviced/utils"
)

const (
	NilPool = "NilPool"
)

var DefaultTestPools = []pool.ResourcePool{
	{
		ID:          "test-pool-id-1",
		CoreLimit:   8,
		MemoryLimit: 0,
	}, {
		ID:          "test-pool-id-2",
		CoreLimit:   4,
		MemoryLimit: 4 * 1024 * 1024 * 1024,
	}, {
		ID:          "test-pool-id-3",
		CoreLimit:   2,
		MemoryLimit: 512 * 1024 * 1024,
	},
}

var DefaultTestHostIPs = []host.HostIPResource{
	{
		HostID:        "test-host-id-1",
		IPAddress:     "127.0.0.1",
		InterfaceName: "test-interface-name-1",
	}, {
		HostID:        "test-host-id-2",
		IPAddress:     "192.168.0.1",
		InterfaceName: "test-interface-name-2",
	}, {
		HostID:        "test-host-id-3",
		IPAddress:     "0.0.0.0",
		InterfaceName: "test-interface-name-3",
	},
}

var (
	ErrNoPoolFound = errors.New("no pool found")
	ErrInvalidPool = errors.New("invalid pool")
)

type PoolAPITest struct {
	api.API
	fail    bool
	pools   []pool.ResourcePool
	hostIPs []host.HostIPResource
}

func DefaultPoolAPI() PoolAPITest {
	return PoolAPITest{pools: DefaultTestPools, hostIPs: DefaultTestHostIPs}
}

func RunCmd(test api.API, args ...string) {
	New(test, utils.TestConfigReader(make(map[string]string))).Run(args)
}

func (t PoolAPITest) GetResourcePools() ([]pool.ResourcePool, error) {
	if t.fail {
		return nil, ErrInvalidPool
	}

	return t.pools, nil
}

func (t PoolAPITest) GetResourcePool(id string) (*pool.ResourcePool, error) {
	if t.fail {
		return nil, ErrInvalidPool
	}

	for _, p := range t.pools {
		if p.ID == id {
			return &p, nil
		}
	}

	return nil, nil
}

func (t PoolAPITest) AddResourcePool(config api.PoolConfig) (*pool.ResourcePool, error) {
	if p, err := t.GetResourcePool(config.PoolID); p != nil || err != nil {
		return nil, ErrInvalidPool
	} else if config.PoolID == NilPool {
		return nil, nil
	}

	p := &pool.ResourcePool{
		ID:          config.PoolID,
		CoreLimit:   config.CoreLimit,
		MemoryLimit: config.MemoryLimit,
	}

	return p, nil
}

func (t PoolAPITest) RemoveResourcePool(id string) error {
	if p, err := t.GetResourcePool(id); err != nil {
		return err
	} else if p == nil {
		return ErrNoPoolFound
	}

	return nil
}

func (t PoolAPITest) GetPoolIPs(id string) (*pool.PoolIPs, error) {
	p, err := t.GetResourcePool(id)
	if err != nil {
		return nil, err
	} else if p == nil {
		return nil, ErrNoPoolFound
	}

	return &pool.PoolIPs{PoolID: p.ID, HostIPs: t.hostIPs}, nil
}

func TestServicedCLI_CmdPoolList_one(t *testing.T) {
	poolID := "test-pool-id-1"

	test := DefaultPoolAPI()
	expected, err := test.GetResourcePool(poolID)
	if err != nil {
		t.Fatal(err)
	}

	var actual pool.ResourcePool
	output := pipeAPI(RunCmd, test, "serviced", "pool", "list", poolID)
	if err := json.Unmarshal(output, &actual); err != nil {
		t.Fatalf("error unmarshalling resource: %s", err)
	}

	// Did you remember to update ResourcePool.Equals?
	if !actual.Equals(expected) {
		t.Fatalf("\ngot:\n%+v\nwant:\n%+v", actual, expected)
	}
}

func TestServicedCLI_CmdPoolList_all(t *testing.T) {
	test := DefaultPoolAPI()
	expected, err := test.GetResourcePools()
	if err != nil {
		t.Fatal(err)
	}

	var actual []*pool.ResourcePool
	output := pipeAPI(RunCmd, test, "serviced", "pool", "list", "--verbose")
	if err := json.Unmarshal(output, &actual); err != nil {
		t.Fatalf("error unmarshalling resource: %s", err)
	}

	// Did you remember to update ResourcePool.Equals?
	if len(actual) != len(expected) {
		t.Fatalf("\ngot:\n%+v\nwant:\n%+v", actual, expected)
	}
	for i, _ := range actual {
		if !actual[i].Equals(&expected[i]) {
			t.Fatalf("\ngot:\n%+v\nwant:\n%+v", actual, expected)
		}
	}
}

func ExampleServicedCLI_CmdPoolList() {
	// Gofmt cleans up the spaces at the end of each row
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "list")
}

func ExampleServicedCLI_CmdPoolList_fail() {
	test := DefaultPoolAPI()
	test.fail = true
	// Error retrieving pool
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list", "test-pool-id-1")
	// Error retrieving all pools
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list")

	// Output:
	// invalid pool
	// invalid pool
}

func ExampleServicedCLI_CmdPoolList_err() {
	test := DefaultPoolAPI()
	test.pools = make([]pool.ResourcePool, 0)
	// Pool not found
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list", "test-pool-id-0")
	// No pools found
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list")

	// Output:
	// pool not found
	// no resource pools found
}

func ExampleServicedCLI_CmdPoolList_complete() {
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "list", "--generate-bash-completion")

	// Output:
	// test-pool-id-1
	// test-pool-id-2
	// test-pool-id-3
}

func ExampleServicedCLI_CmdPoolAdd() {
	// // Bad CoreLimit
	// InitPoolAPITest("serviced", "pool", "add", "test-pool", "abc", "1024", "3")
	// // Bad MemoryLimit
	// InitPoolAPITest("serviced", "pool", "add", "test-pool", "4", "abc", "3")
	// Success
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "add", "test-pool", "3")

	// Output:
	// test-pool
}

func ExampleServicedCLI_CmdPoolAdd_err() {
	pipeAPIStderr(RunCmd, DefaultPoolAPI(), "serviced", "pool", "add", NilPool, "4", "1024", "3")

	// Output:
	// received nil resource pool
}

func ExampleServicedCLI_CmdPoolRemove() {
	pipeAPIStderr(RunCmd, DefaultPoolAPI(), "serviced", "pool", "remove", "test-pool-id-1")

	// Output:
	// test-pool-id-1
}

func ExampleServicedCLI_CmdPoolRemove_usage() {
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "rm")

	// Output:
	// Incorrect Usage.
	//
	// NAME:
	//    remove - Removes an existing resource pool
	//
	// USAGE:
	//    command remove [command options] [arguments...]
	//
	// DESCRIPTION:
	//    serviced pool remove POOLID ...
	//
	// OPTIONS:
}

func ExampleServicedCLI_CmdPoolRemove_err() {
	pipeAPIStderr(RunCmd, DefaultPoolAPI(), "serviced", "pool", "remove", "test-pool-id-0")

	// Output:
	// test-pool-id-0: pool not found
}

func ExampleServicedCLI_CmdPoolRemove_complete() {
	test := DefaultPoolAPI()
	RunCmd(test, "serviced", "pool", "rm", "--generate-bash-completion")
	fmt.Println("")
	RunCmd(test, "serviced", "pool", "rm", "test-pool-id-2", "--generate-bash-completion")

	// Output:
	// test-pool-id-1
	// test-pool-id-2
	// test-pool-id-3
	//
	// test-pool-id-1
	// test-pool-id-3
}

func TestExampleServicedCLI_CmdPoolListIPs(t *testing.T) {
	poolID := "test-pool-id-1"
	test := DefaultPoolAPI()

	var expected []host.HostIPResource
	if ips, err := test.GetPoolIPs(poolID); err != nil {
		t.Fatal(err)
	} else {
		expected = ips.HostIPs
	}

	var actual []host.HostIPResource
	output := pipeAPI(RunCmd, test, "serviced", "pool", "list-ips", poolID, "--verbose")
	if err := json.Unmarshal(output, &actual); err != nil {
		t.Fatalf("error unmarshalling resource: %s", err)
	}

	if !reflect.DeepEqual(actual, expected) {
		t.Fatalf("\ngot:\n%+v\nwant:\n%+v", actual, expected)
	}
}

func ExampleServicedCLI_CmdPoolListIPs() {
	// Gofmt cleans up the spaces at the end of each row
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "list-ips", "test-pool-id-1")
}

func ExampleServicedCLI_CmdPoolListIPs_usage() {
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "list-ips")

	// Output:
	// Incorrect Usage.
	//
	// NAME:
	//    list-ips - Lists the IP addresses for a resource pool
	//
	// USAGE:
	//    command list-ips [command options] [arguments...]
	//
	// DESCRIPTION:
	//    serviced pool list-ips POOLID
	//
	// OPTIONS:
	//    --verbose, -v				Show JSON format
	//    --show-fields 'InterfaceName,IPAddress,Type'	Comma-delimited list describing which fields to display
}

func ExampleServicedCLI_CmdPoolListIPs_fail() {
	pipeAPIStderr(RunCmd, DefaultPoolAPI(), "serviced", "pool", "list-ips", "test-pool-id-0")

	// Output:
	// no pool found
}

func ExampleServicedCLI_CmdPoolListIPs_err() {
	test := DefaultPoolAPI()
	test.hostIPs = nil
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list-ips", "test-pool-id-1")

	// Output:
	// no resource pool IPs found
}
<commit_msg>PoolAPITest maintains internal list of pools<commit_after>// Copyright 2014 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build unit

package cmd

import (
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"testing"

	"github.com/control-center/serviced/cli/api"
	"github.com/control-center/serviced/domain/host"
	"github.com/control-center/serviced/domain/pool"
	"github.com/control-center/serviced/utils"
)

const (
	NilPool = "NilPool"
)

var DefaultTestPools = []pool.ResourcePool{
	{
		ID:          "test-pool-id-1",
		CoreLimit:   8,
		MemoryLimit: 0,
	}, {
		ID:          "test-pool-id-2",
		CoreLimit:   4,
		MemoryLimit: 4 * 1024 * 1024 * 1024,
	}, {
		ID:          "test-pool-id-3",
		CoreLimit:   2,
		MemoryLimit: 512 * 1024 * 1024,
	},
}

var DefaultTestHostIPs = []host.HostIPResource{
	{
		HostID:        "test-host-id-1",
		IPAddress:     "127.0.0.1",
		InterfaceName: "test-interface-name-1",
	}, {
		HostID:        "test-host-id-2",
		IPAddress:     "192.168.0.1",
		InterfaceName: "test-interface-name-2",
	}, {
		HostID:        "test-host-id-3",
		IPAddress:     "0.0.0.0",
		InterfaceName: "test-interface-name-3",
	},
}

var (
	ErrNoPoolFound = errors.New("no pool found")
	ErrInvalidPool = errors.New("invalid pool")
)

type PoolAPITest struct {
	api.API
	fail    bool
	pools   *[]pool.ResourcePool
	hostIPs []host.HostIPResource
}

func DefaultPoolAPI() PoolAPITest {
	test := PoolAPITest{
		pools:   &[]pool.ResourcePool{},
		hostIPs: DefaultTestHostIPs,
	}
	*test.pools = append(*test.pools, DefaultTestPools[:]...)
	return test
}

func RunCmd(test api.API, args ...string) {
	New(test, utils.TestConfigReader(make(map[string]string))).Run(args)
}

func (t PoolAPITest) GetResourcePools() ([]pool.ResourcePool, error) {
	if t.fail {
		return nil, ErrInvalidPool
	}

	return *t.pools, nil
}

func (t PoolAPITest) GetResourcePool(id string) (*pool.ResourcePool, error) {
	if t.fail {
		return nil, ErrInvalidPool
	}

	for _, p := range *t.pools {
		if p.ID == id {
			return &p, nil
		}
	}

	return nil, nil
}

func (t PoolAPITest) AddResourcePool(config api.PoolConfig) (*pool.ResourcePool, error) {
	if p, err := t.GetResourcePool(config.PoolID); p != nil || err != nil {
		return nil, ErrInvalidPool
	} else if config.PoolID == NilPool {
		return nil, nil
	}

	p := &pool.ResourcePool{
		ID:          config.PoolID,
		CoreLimit:   config.CoreLimit,
		MemoryLimit: config.MemoryLimit,
	}

	*t.pools = append(*t.pools, *p)
	return p, nil
}

func (t PoolAPITest) RemoveResourcePool(id string) error {
	if t.fail {
		return ErrInvalidPool
	}

	for i, p := range *t.pools {
		if p.ID == id {
			tmp := *t.pools
			*t.pools = append(tmp[:i], tmp[i+1:]...)
			return nil
		}
	}
	return ErrNoPoolFound
}

func (t PoolAPITest) GetPoolIPs(id string) (*pool.PoolIPs, error) {
	p, err := t.GetResourcePool(id)
	if err != nil {
		return nil, err
	} else if p == nil {
		return nil, ErrNoPoolFound
	}

	return &pool.PoolIPs{PoolID: p.ID, HostIPs: t.hostIPs}, nil
}

func TestServicedCLI_CmdPoolList_one(t *testing.T) {
	poolID := "test-pool-id-1"

	test := DefaultPoolAPI()
	expected, err := test.GetResourcePool(poolID)
	if err != nil {
		t.Fatal(err)
	}

	var actual pool.ResourcePool
	output := pipeAPI(RunCmd, test, "serviced", "pool", "list", poolID)
	if err := json.Unmarshal(output, &actual); err != nil {
		t.Fatalf("error unmarshalling resource: %s", err)
	}

	// Did you remember to update ResourcePool.Equals?
	if !actual.Equals(expected) {
		t.Fatalf("\ngot:\n%+v\nwant:\n%+v", actual, expected)
	}
}

func TestServicedCLI_CmdPoolList_all(t *testing.T) {
	test := DefaultPoolAPI()
	expected, err := test.GetResourcePools()
	if err != nil {
		t.Fatal(err)
	}

	var actual []*pool.ResourcePool
	output := pipeAPI(RunCmd, test, "serviced", "pool", "list", "--verbose")
	if err := json.Unmarshal(output, &actual); err != nil {
		t.Fatalf("error unmarshalling resource: %s", err)
	}

	// Did you remember to update ResourcePool.Equals?
	if len(actual) != len(expected) {
		t.Fatalf("\ngot:\n%+v\nwant:\n%+v", actual, expected)
	}
	for i, _ := range actual {
		if !actual[i].Equals(&expected[i]) {
			t.Fatalf("\ngot:\n%+v\nwant:\n%+v", actual, expected)
		}
	}
}

func ExampleServicedCLI_CmdPoolList() {
	// Gofmt cleans up the spaces at the end of each row
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "list")
}

func ExampleServicedCLI_CmdPoolList_fail() {
	test := DefaultPoolAPI()
	test.fail = true
	// Error retrieving pool
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list", "test-pool-id-1")
	// Error retrieving all pools
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list")

	// Output:
	// invalid pool
	// invalid pool
}

func ExampleServicedCLI_CmdPoolList_err() {
	test := DefaultPoolAPI()
	*test.pools = make([]pool.ResourcePool, 0)

	// Pool not found
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list", "test-pool-id-0")
	// No pools found
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list")

	// Output:
	// pool not found
	// no resource pools found
}

func ExampleServicedCLI_CmdPoolList_complete() {
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "list", "--generate-bash-completion")

	// Output:
	// test-pool-id-1
	// test-pool-id-2
	// test-pool-id-3
}

func ExampleServicedCLI_CmdPoolAdd() {
	// // Bad CoreLimit
	// InitPoolAPITest("serviced", "pool", "add", "test-pool", "abc", "1024", "3")
	// // Bad MemoryLimit
	// InitPoolAPITest("serviced", "pool", "add", "test-pool", "4", "abc", "3")
	// Success
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "add", "test-pool", "3")

	// Output:
	// test-pool
}

func ExampleServicedCLI_CmdPoolAdd_err() {
	pipeAPIStderr(RunCmd, DefaultPoolAPI(), "serviced", "pool", "add", NilPool, "4", "1024", "3")

	// Output:
	// received nil resource pool
}

func ExampleServicedCLI_CmdPoolRemove() {
	pipeAPIStderr(RunCmd, DefaultPoolAPI(), "serviced", "pool", "remove", "test-pool-id-1")

	// Output:
	// test-pool-id-1
}

func ExampleServicedCLI_CmdPoolRemove_usage() {
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "rm")

	// Output:
	// Incorrect Usage.
	//
	// NAME:
	//    remove - Removes an existing resource pool
	//
	// USAGE:
	//    command remove [command options] [arguments...]
	//
	// DESCRIPTION:
	//    serviced pool remove POOLID ...
	//
	// OPTIONS:
}

func ExampleServicedCLI_CmdPoolRemove_err() {
	pipeAPIStderr(RunCmd, DefaultPoolAPI(), "serviced", "pool", "remove", "test-pool-id-0")

	// Output:
	// test-pool-id-0: pool not found
}

func ExampleServicedCLI_CmdPoolRemove_complete() {
	test := DefaultPoolAPI()
	RunCmd(test, "serviced", "pool", "rm", "--generate-bash-completion")
	fmt.Println("")
	RunCmd(test, "serviced", "pool", "rm", "test-pool-id-2", "--generate-bash-completion")

	// Output:
	// test-pool-id-1
	// test-pool-id-2
	// test-pool-id-3
	//
	// test-pool-id-1
	// test-pool-id-3
}

func TestExampleServicedCLI_CmdPoolListIPs(t *testing.T) {
	poolID := "test-pool-id-1"
	test := DefaultPoolAPI()

	var expected []host.HostIPResource
	if ips, err := test.GetPoolIPs(poolID); err != nil {
		t.Fatal(err)
	} else {
		expected = ips.HostIPs
	}

	var actual []host.HostIPResource
	output := pipeAPI(RunCmd, test, "serviced", "pool", "list-ips", poolID, "--verbose")
	if err := json.Unmarshal(output, &actual); err != nil {
		t.Fatalf("error unmarshalling resource: %s", err)
	}

	if !reflect.DeepEqual(actual, expected) {
		t.Fatalf("\ngot:\n%+v\nwant:\n%+v", actual, expected)
	}
}

func ExampleServicedCLI_CmdPoolListIPs() {
	// Gofmt cleans up the spaces at the end of each row
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "list-ips", "test-pool-id-1")
}

func ExampleServicedCLI_CmdPoolListIPs_usage() {
	RunCmd(DefaultPoolAPI(), "serviced", "pool", "list-ips")

	// Output:
	// Incorrect Usage.
	//
	// NAME:
	//    list-ips - Lists the IP addresses for a resource pool
	//
	// USAGE:
	//    command list-ips [command options] [arguments...]
	//
	// DESCRIPTION:
	//    serviced pool list-ips POOLID
	//
	// OPTIONS:
	//    --verbose, -v				Show JSON format
	//    --show-fields 'InterfaceName,IPAddress,Type'	Comma-delimited list describing which fields to display
}

func ExampleServicedCLI_CmdPoolListIPs_fail() {
	pipeAPIStderr(RunCmd, DefaultPoolAPI(), "serviced", "pool", "list-ips", "test-pool-id-0")

	// Output:
	// no pool found
}

func ExampleServicedCLI_CmdPoolListIPs_err() {
	test := DefaultPoolAPI()
	test.hostIPs = nil
	pipeAPIStderr(RunCmd, test, "serviced", "pool", "list-ips", "test-pool-id-1")

	// Output:
	// no resource pool IPs found
}
<|endoftext|>


<commit_before>package cli

// ValueMap is a map fo values with strings for keys
type ValueMap map[string]Value

// Values returns a ValueList of all the values in the map
func (v ValueMap) Values() ValueList {
	var valueList ValueList
	for _, value := range v {
		valueList = append(valueList, value)
	}
	return valueList
}
<commit_msg>add keys method<commit_after>package cli

// ValueMap is a map fo values with strings for keys
type ValueMap map[string]Value

// Keys returns the keys of the value map
func (v ValueMap) Keys() []string {
	var keys []string
	for key := range v {
		keys = append(keys, key)
	}
	return keys
}

// Values returns a ValueList of all the values in the map
func (v ValueMap) Values() ValueList {
	var valueList ValueList
	for _, value := range v {
		valueList = append(valueList, value)
	}
	return valueList
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/velocity-ci/velocity/master/velocity/domain"
	"github.com/velocity-ci/velocity/master/velocity/domain/project/task"
)

func main() {
	version := flag.Bool("v", false, "Show version")
	list := flag.Bool("l", false, "List tasks")

	flag.Parse()

	if *version {
		fmt.Println("Version")
		os.Exit(0)
	} else if *list {
		// look for task ymls and parse them into memory.
		tasks := getTasksFromDirectory("./tasks/")
		// iterate through tasks in memory and list them.
		for _, task := range tasks {
			fmt.Printf("%s: %s (", task.Name, task.Description)
			for _, parameter := range task.Parameters {
				fmt.Printf(" %s= %s ", parameter.Name, parameter.Value)
			}
			fmt.Println(")")
			for _, step := range task.Steps {
				fmt.Printf("\t%s| %s: %s\n", step.GetType(), step.GetDescription(), step.GetDetails())
			}
		}
		os.Exit(0)
	}

	switch os.Args[1] {
	case "run":
		run(os.Args[2])
		break
	}
}

func getTasksFromDirectory(dir string) []domain.Task {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		panic(err)
	}

	tasks := []domain.Task{}

	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".yml") || strings.HasSuffix(file.Name(), ".yaml") {
			taskYml, _ := ioutil.ReadFile(fmt.Sprintf("%s%s", dir, file.Name()))
			task := task.ResolveTaskFromYAML(string(taskYml))
			tasks = append(tasks, task)
		}
	}

	return tasks
}

func run(taskName string) {
	tasks := getTasksFromDirectory("./tasks/")

	var task *domain.Task
	// find Task requested
	for _, t := range tasks {
		if t.Name == taskName {
			task = &t
			break
		}
	}

	if task == nil {
		panic(fmt.Sprintf("Task %s not found\n%s", taskName, tasks))
	}

	fmt.Printf("Running task: %s (from: %s)\n", task.Name, taskName)

	// Resolve parameters
	reader := bufio.NewReader(os.Stdin)
	resolvedParams := []domain.Parameter{}
	for _, p := range task.Parameters {
		// get real value for parameter (ask or from env)
		inputText := ""
		for len(strings.TrimSpace(inputText)) < 1 {
			fmt.Printf("Enter a value for %s (default: %s): ", p.Name, p.Value)
			inputText, _ = reader.ReadString('\n')
		}
		p.Value = strings.TrimSpace(inputText)
		resolvedParams = append(resolvedParams, p)
	}
	task.Parameters = resolvedParams
	task.UpdateParams()
	task.SetEmitter(func(s string) { fmt.Printf("    %s\n", s) })

	// Run each step unless they fail (optional)
	for _, step := range task.Steps {
		step.Execute()
	}
}
<commit_msg>[cli] Getting tasks from subdirectories<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/velocity-ci/velocity/master/velocity/domain"
	"github.com/velocity-ci/velocity/master/velocity/domain/project/task"
)

func main() {
	version := flag.Bool("v", false, "Show version")
	list := flag.Bool("l", false, "List tasks")

	flag.Parse()

	if *version {
		fmt.Println("Version")
		os.Exit(0)
	} else if *list {
		// look for task ymls and parse them into memory.
		tasks := getTasksFromDirectory("./tasks/")
		// iterate through tasks in memory and list them.
		for _, task := range tasks {
			fmt.Printf("%s: %s (", task.Name, task.Description)
			for _, parameter := range task.Parameters {
				fmt.Printf(" %s= %s ", parameter.Name, parameter.Value)
			}
			fmt.Println(")")
			for _, step := range task.Steps {
				fmt.Printf("\t%s| %s: %s\n", step.GetType(), step.GetDescription(), step.GetDetails())
			}
		}
		os.Exit(0)
	}

	switch os.Args[1] {
	case "run":
		run(os.Args[2])
		break
	}
}

func getTasksFromDirectory(dir string) []domain.Task {
	tasks := []domain.Task{}

	filepath.Walk(dir, func(path string, f os.FileInfo, err error) error {
		if !f.IsDir() && strings.HasSuffix(f.Name(), ".yml") || strings.HasSuffix(f.Name(), ".yaml") {
			taskYml, _ := ioutil.ReadFile(fmt.Sprintf("%s%s", dir, f.Name()))
			task := task.ResolveTaskFromYAML(string(taskYml))
			tasks = append(tasks, task)
		}
		return nil
	})

	return tasks
}

func run(taskName string) {
	tasks := getTasksFromDirectory("./tasks/")

	var task *domain.Task
	// find Task requested
	for _, t := range tasks {
		if t.Name == taskName {
			task = &t
			break
		}
	}

	if task == nil {
		panic(fmt.Sprintf("Task %s not found\n%s", taskName, tasks))
	}

	fmt.Printf("Running task: %s (from: %s)\n", task.Name, taskName)

	// Resolve parameters
	reader := bufio.NewReader(os.Stdin)
	resolvedParams := []domain.Parameter{}
	for _, p := range task.Parameters {
		// get real value for parameter (ask or from env)
		inputText := ""
		for len(strings.TrimSpace(inputText)) < 1 {
			fmt.Printf("Enter a value for %s (default: %s): ", p.Name, p.Value)
			inputText, _ = reader.ReadString('\n')
		}
		p.Value = strings.TrimSpace(inputText)
		resolvedParams = append(resolvedParams, p)
	}
	task.Parameters = resolvedParams
	task.UpdateParams()
	task.SetEmitter(func(s string) { fmt.Printf("    %s\n", s) })

	// Run each step unless they fail (optional)
	for _, step := range task.Steps {
		step.Execute()
	}
}
<|endoftext|>


<commit_before>package git

import (
	"compress/zlib"
	"crypto/sha1"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

type Repository struct {
	Path       string
	Bare       bool
	root       string
	pack       *Pack
	packedRefs *PackedRefs
}

func Open(path string) (*Repository, error) {
	path = filepath.Clean(path)
	fi, err := os.Stat(path)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, fmt.Errorf("Not a git repository: %s", path)
	}

	repo := &Repository{
		Path: path,
		root: path,
	}
	defer func() {
		if repo != nil {
			repo.openPackedRefs()
		}
	}()
	if strings.HasSuffix(path, ".git") {
		repo.Bare = true
		return repo, nil
	}

	files, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, err
	}
	for _, file := range files {
		if file.Name() == ".git" {
			repo.root = filepath.Join(repo.root, ".git")
			return repo, nil
		}
	}
	return nil, fmt.Errorf("Not a git repository: %s", path)
}

func (r *Repository) Object(id SHA1) (Object, error) {
	return r.readObject(id, nil, false)
}

func (r *Repository) Resolve(obj Object) error {
	if obj.Resolved() || obj.SHA1().Empty() {
		return nil
	}
	_, err := r.readObject(obj.SHA1(), obj, false)
	return err
}

func (r *Repository) readObject(id SHA1, obj Object, headerOnly bool) (Object, error) {
	entry, err := r.entry(id)
	if err != nil {
		return nil, err
	}
	defer entry.Close()

	if obj == nil {
		obj = newObject(entry.Type(), id, r)
	}

	if headerOnly {
		return obj, nil
	}
	b, err := entry.ReadAll()
	if err != nil {
		return nil, err
	}
	err = obj.Parse(b)
	return obj, err
}

func (r *Repository) entry(id SHA1) (objectEntry, error) {
	if r.pack == nil {
		if err := r.openPack(); err != nil {
			return nil, err
		}
	}

	if entry, err := r.pack.entry(id); err == nil {
		return entry, err
	}
	return newLooseObjectEntry(r.root, id)
}

func (r *Repository) openPack() error {
	pattern := filepath.Join(r.root, "objects", "pack", "pack-*.pack")
	files, err := filepath.Glob(pattern)
	if err != nil {
		return err
	}
	switch len(files) {
	case 0: // set empty pack
		r.pack = &Pack{idx: &PackIndexV2{}}
	case 1:
		pack, err := OpenPack(files[0])
		if err != nil {
			return err
		}
		r.pack = pack
	default:
		return errors.New("Found more than 1 pack file")
	}
	return nil
}

func (r *Repository) writeObject(typ string, data ObjectData) (id SHA1, err error) {
	var path string
	defer func() {
		if err != nil && path != "" {
			os.Remove(path)
		}
	}()
	if id, path, err = r.writeObjectData(typ, data); err != nil {
		return
	}
	err = r.storeAsObject(path, id)
	return
}

func (r *Repository) writeObjectData(typ string, data ObjectData) (id SHA1, path string, err error) {
	var f *os.File
	if f, err = ioutil.TempFile("", "go-git"); err != nil {
		return
	}
	defer f.Close()
	path = f.Name()

	zw := zlib.NewWriter(f)
	defer zw.Close()
	hash := sha1.New()
	w := io.MultiWriter(hash, zw)

	if _, err = fmt.Fprintf(w, "%s %d%c", typ, data.Size(), 0); err != nil {
		return
	}
	if _, err = io.Copy(w, data); err != nil {
		return
	}
	id = SHA1FromBytes(hash.Sum(nil))
	return
}

func (r *Repository) storeAsObject(path string, id SHA1) error {
	s := id.String()
	dir := filepath.Join(r.root, "objects", s[:2])
	if err := os.MkdirAll(dir, 0777); err != nil {
		return err
	}
	return os.Rename(path, filepath.Join(dir, s[2:]))
}
<commit_msg>Support multiple pack files<commit_after>package git

import (
	"compress/zlib"
	"crypto/sha1"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

type Repository struct {
	Path       string
	Bare       bool
	root       string
	packs      []*Pack
	packedRefs *PackedRefs
}

func Open(path string) (*Repository, error) {
	path = filepath.Clean(path)
	fi, err := os.Stat(path)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, fmt.Errorf("Not a git repository: %s", path)
	}

	repo := &Repository{
		Path: path,
		root: path,
	}
	defer func() {
		if repo != nil {
			repo.openPackedRefs()
		}
	}()
	if strings.HasSuffix(path, ".git") {
		repo.Bare = true
		return repo, nil
	}

	files, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, err
	}
	for _, file := range files {
		if file.Name() == ".git" {
			repo.root = filepath.Join(repo.root, ".git")
			return repo, nil
		}
	}
	return nil, fmt.Errorf("Not a git repository: %s", path)
}

func (r *Repository) Object(id SHA1) (Object, error) {
	return r.readObject(id, nil, false)
}

func (r *Repository) Resolve(obj Object) error {
	if obj.Resolved() || obj.SHA1().Empty() {
		return nil
	}
	_, err := r.readObject(obj.SHA1(), obj, false)
	return err
}

func (r *Repository) readObject(id SHA1, obj Object, headerOnly bool) (Object, error) {
	entry, err := r.entry(id)
	if err != nil {
		return nil, err
	}
	defer entry.Close()

	if obj == nil {
		obj = newObject(entry.Type(), id, r)
	}

	if headerOnly {
		return obj, nil
	}
	b, err := entry.ReadAll()
	if err != nil {
		return nil, err
	}
	err = obj.Parse(b)
	return obj, err
}

func (r *Repository) entry(id SHA1) (objectEntry, error) {
	if r.packs == nil {
		if err := r.openPack(); err != nil {
			return nil, err
		}
	}

	for _, pack := range r.packs {
		if entry, err := pack.entry(id); err == nil {
			return entry, err
		}
	}
	return newLooseObjectEntry(r.root, id)
}

func (r *Repository) openPack() error {
	pattern := filepath.Join(r.root, "objects", "pack", "pack-*.pack")
	files, err := filepath.Glob(pattern)
	if err != nil {
		return err
	}
	packs := []*Pack{}
	for _, file := range files {
		pack, err := OpenPack(file)
		if err != nil {
			return err
		}
		packs = append(packs, pack)
	}
	r.packs = packs
	return nil
}

func (r *Repository) writeObject(typ string, data ObjectData) (id SHA1, err error) {
	var path string
	defer func() {
		if err != nil && path != "" {
			os.Remove(path)
		}
	}()
	if id, path, err = r.writeObjectData(typ, data); err != nil {
		return
	}
	err = r.storeAsObject(path, id)
	return
}

func (r *Repository) writeObjectData(typ string, data ObjectData) (id SHA1, path string, err error) {
	var f *os.File
	if f, err = ioutil.TempFile("", "go-git"); err != nil {
		return
	}
	defer f.Close()
	path = f.Name()

	zw := zlib.NewWriter(f)
	defer zw.Close()
	hash := sha1.New()
	w := io.MultiWriter(hash, zw)

	if _, err = fmt.Fprintf(w, "%s %d%c", typ, data.Size(), 0); err != nil {
		return
	}
	if _, err = io.Copy(w, data); err != nil {
		return
	}
	id = SHA1FromBytes(hash.Sum(nil))
	return
}

func (r *Repository) storeAsObject(path string, id SHA1) error {
	s := id.String()
	dir := filepath.Join(r.root, "objects", s[:2])
	if err := os.MkdirAll(dir, 0777); err != nil {
		return err
	}
	return os.Rename(path, filepath.Join(dir, s[2:]))
}
<|endoftext|>


<commit_before>package native

import (
	"sync"

	"github.com/CyCoreSystems/ari"

	"golang.org/x/net/context"
)

// busChannelBuffer defines the buffer size of the subscription channels
var busChannelBuffer = 100

// bus is an event bus for ARI events.  It receives and
// redistributes events based on a subscription
// model.
type bus struct {
	subs []*subscription // The list of subscriptions

	mu sync.Mutex
}

func (b *bus) addSubscription(s *subscription) {
	b.mu.Lock()
	b.subs = append(b.subs, s)
	b.mu.Unlock()
}

func (b *bus) removeSubscription(s *subscription) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for i, si := range b.subs {
		if s == si {
			// Subs are pointers, so we have to explicitly remove them
			// to prevent memory leaks
			b.subs[i] = b.subs[len(b.subs)-1] // replace the current with the end
			b.subs[len(b.subs)-1] = nil       // remove the end
			b.subs = b.subs[:len(b.subs)-1]   // lop off the end
			return
		}
	}
}

// Send sends the message on the bus
func (b *bus) Send(msg *ari.Message) {
	b.send(msg)
}

func (b *bus) send(msg *ari.Message) {
	e := ari.Events.Parse(msg)

	Logger.Debug("Received event", "event", e)

	// Disseminate the message to the subscribers
	b.mu.Lock()
	defer b.mu.Unlock()
	for _, s := range b.subs {
		for _, topic := range s.events {
			if topic == e.GetType() || topic == ari.Events.All {
				select {
				case s.C <- e:
				default: // never block
				}
			}
		}
	}
}

// startbus creates and returns the event bus.
func startbus(ctx context.Context) *bus {
	b := &bus{
		subs: []*subscription{},
	}

	// Listen for stop and shut down subscriptions, as required
	go func() {
		<-ctx.Done()
		b.Stop()
		return
	}()

	return b
}

// Stop the bus.  Cancels all subscriptions
// and stops listening for events.
func (b *bus) Stop() {
	// Close all subscriptions
	b.mu.Lock()
	if b.subs != nil {
		for i, s := range b.subs {
			s.closeChan()
			b.subs[i] = nil
		}
		b.subs = nil
	}
	b.mu.Unlock()
}

// A Subscription is a wrapped channel for receiving
// events from the ARI event bus.
type subscription struct {
	b      *bus           // reference to the event bus
	events []string       // list of events to listen for
	C      chan ari.Event // channel for sending events to the subscriber
	mu     sync.Mutex
	Closed bool
}

// newSubscription creates a new, unattached subscription
func newSubscription(eTypes ...string) *subscription {
	return &subscription{
		events: eTypes,
		C:      make(chan ari.Event, busChannelBuffer),
	}
}

// Subscribe returns a subscription to the given list
// of event types
func (b *bus) Subscribe(eTypes ...string) ari.Subscription {
	return b.subscribe(eTypes...)
}

// subscribe returns a subscription to the given list
// of event types
func (b *bus) subscribe(eTypes ...string) *subscription {
	s := newSubscription(eTypes...)
	s.b = b
	b.addSubscription(s)
	return s
}

// Events returns the events channel
func (s *subscription) Events() chan ari.Event {
	return s.C
}

// Next blocks for the next event in the subscription,
// returning that event when it arrives or nil if
// the subscription is canceled.
// Normally, one would listen to subscription.C directly,
// but this is a convenience function for providing a
// context to alternately cancel.
func (s *subscription) Next(ctx context.Context) ari.Event {
	select {
	case <-ctx.Done():
		return nil
	case e := <-s.C:
		return e
	}
}

func (s *subscription) closeChan() {
	s.mu.Lock()
	if s.C != nil {
		close(s.C)
		s.C = nil
	}
	s.Closed = true
	s.mu.Unlock()
}

// Cancel cancels the subscription and removes it from
// the event bus.
func (s *subscription) Cancel() {
	if s.b != nil {
		s.b.removeSubscription(s)
	}
	s.closeChan()
}

// Once listens for the first event of the provided types,
// returning a channel which supplies that event.
func (b *bus) Once(ctx context.Context, eTypes ...string) <-chan ari.Event {
	s := b.subscribe(eTypes...)

	ret := make(chan ari.Event, busChannelBuffer)

	// Stop subscription after one event
	go func() {
		select {
		case ret <- <-s.C:
		case <-ctx.Done():
		}
		close(ret)
		s.Cancel()
	}()
	return ret
}
<commit_msg>v3 - native - remove noisy event logging<commit_after>package native

import (
	"sync"

	"github.com/CyCoreSystems/ari"

	"golang.org/x/net/context"
)

// busChannelBuffer defines the buffer size of the subscription channels
var busChannelBuffer = 100

// bus is an event bus for ARI events.  It receives and
// redistributes events based on a subscription
// model.
type bus struct {
	subs []*subscription // The list of subscriptions

	mu sync.Mutex
}

func (b *bus) addSubscription(s *subscription) {
	b.mu.Lock()
	b.subs = append(b.subs, s)
	b.mu.Unlock()
}

func (b *bus) removeSubscription(s *subscription) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for i, si := range b.subs {
		if s == si {
			// Subs are pointers, so we have to explicitly remove them
			// to prevent memory leaks
			b.subs[i] = b.subs[len(b.subs)-1] // replace the current with the end
			b.subs[len(b.subs)-1] = nil       // remove the end
			b.subs = b.subs[:len(b.subs)-1]   // lop off the end
			return
		}
	}
}

// Send sends the message on the bus
func (b *bus) Send(msg *ari.Message) {
	b.send(msg)
}

func (b *bus) send(msg *ari.Message) {
	e := ari.Events.Parse(msg)

	//	Logger.Debug("Received event", "event", e)

	// Disseminate the message to the subscribers
	b.mu.Lock()
	defer b.mu.Unlock()
	for _, s := range b.subs {
		for _, topic := range s.events {
			if topic == e.GetType() || topic == ari.Events.All {
				select {
				case s.C <- e:
				default: // never block
				}
			}
		}
	}
}

// startbus creates and returns the event bus.
func startbus(ctx context.Context) *bus {
	b := &bus{
		subs: []*subscription{},
	}

	// Listen for stop and shut down subscriptions, as required
	go func() {
		<-ctx.Done()
		b.Stop()
		return
	}()

	return b
}

// Stop the bus.  Cancels all subscriptions
// and stops listening for events.
func (b *bus) Stop() {
	// Close all subscriptions
	b.mu.Lock()
	if b.subs != nil {
		for i, s := range b.subs {
			s.closeChan()
			b.subs[i] = nil
		}
		b.subs = nil
	}
	b.mu.Unlock()
}

// A Subscription is a wrapped channel for receiving
// events from the ARI event bus.
type subscription struct {
	b      *bus           // reference to the event bus
	events []string       // list of events to listen for
	C      chan ari.Event // channel for sending events to the subscriber
	mu     sync.Mutex
	Closed bool
}

// newSubscription creates a new, unattached subscription
func newSubscription(eTypes ...string) *subscription {
	return &subscription{
		events: eTypes,
		C:      make(chan ari.Event, busChannelBuffer),
	}
}

// Subscribe returns a subscription to the given list
// of event types
func (b *bus) Subscribe(eTypes ...string) ari.Subscription {
	return b.subscribe(eTypes...)
}

// subscribe returns a subscription to the given list
// of event types
func (b *bus) subscribe(eTypes ...string) *subscription {
	s := newSubscription(eTypes...)
	s.b = b
	b.addSubscription(s)
	return s
}

// Events returns the events channel
func (s *subscription) Events() chan ari.Event {
	return s.C
}

// Next blocks for the next event in the subscription,
// returning that event when it arrives or nil if
// the subscription is canceled.
// Normally, one would listen to subscription.C directly,
// but this is a convenience function for providing a
// context to alternately cancel.
func (s *subscription) Next(ctx context.Context) ari.Event {
	select {
	case <-ctx.Done():
		return nil
	case e := <-s.C:
		return e
	}
}

func (s *subscription) closeChan() {
	s.mu.Lock()
	if s.C != nil {
		close(s.C)
		s.C = nil
	}
	s.Closed = true
	s.mu.Unlock()
}

// Cancel cancels the subscription and removes it from
// the event bus.
func (s *subscription) Cancel() {
	if s.b != nil {
		s.b.removeSubscription(s)
	}
	s.closeChan()
}

// Once listens for the first event of the provided types,
// returning a channel which supplies that event.
func (b *bus) Once(ctx context.Context, eTypes ...string) <-chan ari.Event {
	s := b.subscribe(eTypes...)

	ret := make(chan ari.Event, busChannelBuffer)

	// Stop subscription after one event
	go func() {
		select {
		case ret <- <-s.C:
		case <-ctx.Done():
		}
		close(ret)
		s.Cancel()
	}()
	return ret
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/fujiwara/fluent-agent-hydra/fluent"
	"github.com/fujiwara/fluent-agent-hydra/hydra"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
)

const (
	defaultMessageKey = "message"
)

func main() {
	var (
		configFile string
		help       bool
		fieldName  string
	)
	flag.StringVar(&configFile, "c", "", "configuration file path")
	flag.BoolVar(&help, "h", false, "show help message")
	flag.BoolVar(&help, "help", false, "show help message")
	flag.StringVar(&fieldName, "f", defaultMessageKey, "fieldname of fluentd log message attribute (DEFAULT: message)")
	flag.Parse()

	if help {
		usage()
	}
	if pprofile := os.Getenv("PPROF"); pprofile != "" {
		f, err := os.Create(pprofile)
		if err != nil {
			log.Fatal("Can't create profiling stat file.", err)
		}
		log.Println("profiling stat file", f.Name())
		pprof.StartCPUProfile(f)
	}

	done := make(chan os.Signal, 1)
	signal.Notify(done, os.Interrupt)

	if configFile != "" {
		config, err := hydra.ReadConfig(configFile)
		if err != nil {
			log.Println("Can't load config", err)
			os.Exit(2)
		}
		runWithConfig(config)
	} else if args := flag.Args(); len(args) >= 3 {
		config := newConfig(args, fieldName)
		runWithConfig(config)
	} else {
		usage()
	}

	sig := <-done
	log.Println("SIGNAL", sig, "exit")
	pprof.StopCPUProfile()
	os.Exit(0)
}

func usage() {
	fmt.Println("Usage of fluent-agent-hydra")
	fmt.Println("")
	fmt.Println("  fluent-agent-hydra -c config.toml")
	fmt.Println("  fluent-agent-hydra [options] TAG TARGET_FILE PRIMARY_SERVER SECONDARY_SERVER")
	fmt.Println("")
	flag.PrintDefaults()
	os.Exit(1)
}

func newConfig(args []string, fieldName string) hydra.Config {
	tag := args[0]
	file := args[1]
	servers := args[2:]

	logs := make([]hydra.ConfigLogfile, 1)
	logs[0] = hydra.ConfigLogfile{Tag: tag, File: file}
	return hydra.Config{
		Servers:   servers,
		Logs:      logs,
		FieldName: fieldName,
	}
}

func runWithConfig(config hydra.Config) {
	loggers := make([]*fluent.Fluent, len(config.Servers))
	for i, server := range config.Servers {
		logger, err := fluent.New(fluent.Config{Server: server})
		if err != nil {
			log.Println("[warning] Can't initialize fluentd server.", err)
		} else {
			log.Println("[info] server", server)
		}
		loggers[i] = logger
	}

	ch := hydra.NewChannel()
	for _, logfile := range config.Logs {
		var tag string
		if config.TagPrefix != "" {
			tag = config.TagPrefix + "." + logfile.Tag
		} else {
			tag = logfile.Tag
		}
		go hydra.Trail(logfile.File, tag, ch)
	}

	var fieldName string
	if config.FieldName != "" {
		fieldName = config.FieldName
	} else {
		fieldName = defaultMessageKey
	}

	go hydra.Forward(ch, fieldName, loggers...)
}
<commit_msg>trap signals INT, HUP, TERM, QUIT<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/fujiwara/fluent-agent-hydra/fluent"
	"github.com/fujiwara/fluent-agent-hydra/hydra"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"syscall"
)

const (
	defaultMessageKey = "message"
)

var (
	trapSignals = []os.Signal{
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT}
)

func main() {
	var (
		configFile string
		help       bool
		fieldName  string
	)
	flag.StringVar(&configFile, "c", "", "configuration file path")
	flag.BoolVar(&help, "h", false, "show help message")
	flag.BoolVar(&help, "help", false, "show help message")
	flag.StringVar(&fieldName, "f", defaultMessageKey, "fieldname of fluentd log message attribute (DEFAULT: message)")
	flag.Parse()

	if help {
		usage()
	}
	if pprofile := os.Getenv("PPROF"); pprofile != "" {
		f, err := os.Create(pprofile)
		if err != nil {
			log.Fatal("[error] Can't create profiling stat file.", err)
		}
		log.Println("[info] StartCPUProfile() stat file", f.Name())
		pprof.StartCPUProfile(f)
	}

	done := make(chan os.Signal, 1)
	signal.Notify(done, trapSignals...)
	if configFile != "" {
		config, err := hydra.ReadConfig(configFile)
		if err != nil {
			log.Println("Can't load config", err)
			os.Exit(2)
		}
		runWithConfig(config)
	} else if args := flag.Args(); len(args) >= 3 {
		config := newConfig(args, fieldName)
		runWithConfig(config)
	} else {
		usage()
	}
	sig := <-done
	log.Println("[info] SIGNAL", sig, ", exit.")
	pprof.StopCPUProfile()
	os.Exit(0)
}

func usage() {
	fmt.Println("Usage of fluent-agent-hydra")
	fmt.Println("")
	fmt.Println("  fluent-agent-hydra -c config.toml")
	fmt.Println("  fluent-agent-hydra [options] TAG TARGET_FILE PRIMARY_SERVER SECONDARY_SERVER")
	fmt.Println("")
	flag.PrintDefaults()
	os.Exit(1)
}

func newConfig(args []string, fieldName string) hydra.Config {
	tag := args[0]
	file := args[1]
	servers := args[2:]

	logs := make([]hydra.ConfigLogfile, 1)
	logs[0] = hydra.ConfigLogfile{Tag: tag, File: file}
	return hydra.Config{
		Servers:   servers,
		Logs:      logs,
		FieldName: fieldName,
	}
}

func runWithConfig(config hydra.Config) {
	loggers := make([]*fluent.Fluent, len(config.Servers))
	for i, server := range config.Servers {
		logger, err := fluent.New(fluent.Config{Server: server})
		if err != nil {
			log.Println("[warning] Can't initialize fluentd server.", err)
		} else {
			log.Println("[info] server", server)
		}
		loggers[i] = logger
	}

	ch := hydra.NewChannel()
	for _, logfile := range config.Logs {
		var tag string
		if config.TagPrefix != "" {
			tag = config.TagPrefix + "." + logfile.Tag
		} else {
			tag = logfile.Tag
		}
		go hydra.Trail(logfile.File, tag, ch)
	}

	var fieldName string
	if config.FieldName != "" {
		fieldName = config.FieldName
	} else {
		fieldName = defaultMessageKey
	}

	go hydra.Forward(ch, fieldName, loggers...)
}
<|endoftext|>


<commit_before>package sphere

import (
	"fmt"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"github.com/rs/xid"
)

const (
	// Read buffer size for websocket upgrader
	readBufferSize = 1024
	// Write buffer size for websocker upgrader
	writeBufferSize = 1024
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second
	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second
	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
	// Maximum message size allowed from peer.
	maxMessageSize = 512
	// TextMessage denotes a text data message. The text message payload is
	// interpreted as UTF-8 encoded text data.
	TextMessage = 1
	// BinaryMessage denotes a binary data message.
	BinaryMessage = 2
	// CloseMessage denotes a close control message. The optional message
	// payload contains a numeric code and text. Use the FormatCloseMessage
	// function to format a close message payload.
	CloseMessage = 8
	// PingMessage denotes a ping control message. The optional message payload
	// is UTF-8 encoded text.
	PingMessage = 9
	// PongMessage denotes a ping control message. The optional message payload
	// is UTF-8 encoded text.
	PongMessage = 10
	// Message Types
)

var (
	// Websocket Upgrader
	upgrader = websocket.Upgrader{
		ReadBufferSize:  readBufferSize,
		WriteBufferSize: writeBufferSize,
	}
	// Guid to generate globally unique id
	guid = xid.New()
)

// NewSphere creates a new instance of Sphere
func NewSphere(brokers ...Agent) *Sphere {
	// declare agent
	var broker Agent
	// set declared agent if parameter exists
	for _, i := range brokers {
		broker = i
		break
	}
	if broker == nil {
		broker = NewSimpleBroker()
	}
	// creates sphere instance
	sphere := &Sphere{
		agent:       broker,
		connections: NewConnectionMap(),
		channels:    NewChannelMap(),
		models:      NewChannelModelMap(),
	}
	return sphere
}

// Sphere represents an entire Websocket instance
type Sphere struct {
	// a broker agent
	agent Agent
	// list of active connections
	connections ConnectionMap
	// list of channels
	channels ChannelMap
	// list of models
	models ChannelModelMap
}

// Handler handles and creates websocket connection
func (sphere *Sphere) Handler(w http.ResponseWriter, r *http.Request) {
	if conn, err := NewConnection(w, r); err == nil {
		sphere.connections.Set(conn.id, conn)
		go sphere.write(conn)
		sphere.read(conn)
		defer func() {
			for item := range conn.channels.Iter() {
				channel := item.Val
				sphere.unsubscribe(channel.namespace, channel.room, conn)
			}
			sphere.connections.Remove(conn.id)
		}()
	}
}

// ChannelModels channel models
func (sphere *Sphere) ChannelModels(models ...IChannels) {
	for _, model := range models {
		if !sphere.models.Has(model.Namespace()) {
			sphere.models.Set(model.Namespace(), model)
		} else {
			panic(fmt.Sprintf("model \"%s\" is already existed", model.Namespace()))
		}
	}
}

func (sphere *Sphere) write(conn *Connection) {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
	}()
	for {
		select {
		case msg, ok := <-conn.send:
			if !ok {
				conn.emit(websocket.CloseMessage, []byte{})
				return
			}
			if err := conn.emit(websocket.TextMessage, msg); err != nil {
				return
			}
		case <-ticker.C:
			if err := conn.emit(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (sphere *Sphere) read(conn *Connection) {
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			return
		}
		if msg != nil {
			go sphere.process(conn, msg)
		}
	}
}

func (sphere *Sphere) process(conn *Connection, msg []byte) {
	p, err := ParsePacket(msg)
	if err != nil {
		fmt.Printf("Error: %v - %v", err.Error(), string(msg[:]))
		return
	}
	if p != nil {
		switch p.Type {
		case PacketTypeChannel:
			if p.Namespace != "" && p.Room != "" {
				p.Machine = sphere.agent.ID()
				sphere.publish(p, conn)
			} else {
				p.Error = ErrBadScheme
				if json, err := p.toJSON(); err == nil {
					conn.send <- json
				}
			}
		case PacketTypeSubscribe:
			if p.Namespace != "" && p.Room != "" {
				err := sphere.subscribe(p.Namespace, p.Room, conn)
				r := p.Response()
				r.SetError(err)
				conn.emit(TextMessage, r, true)
			} else {
				p.Error = ErrBadScheme
				if json, err := p.toJSON(); err == nil {
					conn.send <- json
				}
			}
		case PacketTypeUnsubscribe:
			if p.Namespace != "" && p.Room != "" {
				if sphere.models.Has(p.Namespace) {
					sphere.unsubscribe(p.Namespace, p.Room, conn)
				} else {
					p.Error = ErrUnsupportedNamespace
					if json, err := p.toJSON(); err == nil {
						conn.send <- json
					}
				}
			} else {
				p.Error = ErrBadScheme
				if json, err := p.toJSON(); err == nil {
					conn.send <- json
				}
			}
		case PacketTypePing:
			r := p.Response()
			if json, err := r.toJSON(); err == nil {
				conn.send <- json
			}
		case PacketTypeMessage:
			conn.receive <- p
		}
	}
}

func (sphere *Sphere) channel(namespace string, room string, autoCreateOpts ...bool) *Channel {
	c := make(chan *Channel)
	autoCreateOpt := false
	for _, opt := range autoCreateOpts {
		autoCreateOpt = opt
		break
	}
	name := sphere.agent.ChannelName(namespace, room)
	go func() {
		if tmp, ok := sphere.channels.Get(name); ok {
			c <- tmp
		} else {
			if autoCreateOpt {
				channel := NewChannel(namespace, room)
				sphere.channels.Set(name, channel)
				c <- channel
			} else {
				c <- nil
			}
		}
	}()
	return <-c
}

func (sphere *Sphere) subscribe(namespace string, room string, conn *Connection) error {
	var model IChannels
	if !sphere.models.Has(namespace) {
		return ErrUnsupportedNamespace
	}
	if tmp, ok := sphere.models.Get(namespace); ok {
		model = tmp
	} else {
		return ErrUnsupportedNamespace
	}
	if accept := model.Subscribe(room, conn); !accept {
		return ErrNotAuthorized
	}
	channel := sphere.channel(namespace, room, true)
	if channel == nil {
		return ErrBadStatus
	}
	if err := channel.subscribe(conn); err == nil && !sphere.agent.IsSubscribed(channel.namespace, channel.room) {
		go sphere.agent.OnSubscribe(channel)
	}
	return nil
}

func (sphere *Sphere) unsubscribe(namespace string, room string, conn *Connection) error {
	var model IChannels
	if !sphere.models.Has(namespace) {
		return ErrUnsupportedNamespace
	}
	if tmp, ok := sphere.models.Get(namespace); ok {
		model = tmp
	} else {
		return ErrUnsupportedNamespace
	}
	if success := model.Disconnect(room, conn); !success {
		return ErrBadStatus
	}
	channel := sphere.channel(namespace, room, false)
	if channel == nil {
		return ErrBadStatus
	}
	if err := channel.unsubscribe(conn); err == nil && channel.connections.Count() == 0 {
		if sphere.agent.IsSubscribed(channel.namespace, channel.room) {
			go sphere.agent.OnUnsubscribe(channel)
		}
	}
	return nil
}

func (sphere *Sphere) publish(p *Packet, conn *Connection) error {
	var model IChannels
	if !sphere.models.Has(p.Namespace) {
		return ErrUnsupportedNamespace
	}
	if tmp, ok := sphere.models.Get(p.Namespace); ok {
		model = tmp
	} else {
		return ErrUnsupportedNamespace
	}
	channel := sphere.channel(p.Namespace, p.Room)
	if channel == nil {
		return ErrBadStatus
	}
	if !channel.isSubscribed(conn) {
		return ErrNotAuthorized
	}
	msg := p.Message
	if msg == nil || msg.Event == "" {
		return ErrBadScheme
	}
	res, err := model.Receive(msg.Event, msg.Data)
	if err != nil {
		return err
	}
	d := p.Response()
	if str, ok := res.(string); ok {
		d.Message.Data = str
	}
	if sphere.agent.IsSubscribed(channel.namespace, channel.room) {
		sphere.agent.OnPublish(channel, d)
	}
	return nil
}
<commit_msg>data is string, no need to cast type<commit_after>package sphere

import (
	"fmt"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"github.com/rs/xid"
)

const (
	// Read buffer size for websocket upgrader
	readBufferSize = 1024
	// Write buffer size for websocker upgrader
	writeBufferSize = 1024
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second
	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second
	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
	// Maximum message size allowed from peer.
	maxMessageSize = 512
	// TextMessage denotes a text data message. The text message payload is
	// interpreted as UTF-8 encoded text data.
	TextMessage = 1
	// BinaryMessage denotes a binary data message.
	BinaryMessage = 2
	// CloseMessage denotes a close control message. The optional message
	// payload contains a numeric code and text. Use the FormatCloseMessage
	// function to format a close message payload.
	CloseMessage = 8
	// PingMessage denotes a ping control message. The optional message payload
	// is UTF-8 encoded text.
	PingMessage = 9
	// PongMessage denotes a ping control message. The optional message payload
	// is UTF-8 encoded text.
	PongMessage = 10
	// Message Types
)

var (
	// Websocket Upgrader
	upgrader = websocket.Upgrader{
		ReadBufferSize:  readBufferSize,
		WriteBufferSize: writeBufferSize,
	}
	// Guid to generate globally unique id
	guid = xid.New()
)

// NewSphere creates a new instance of Sphere
func NewSphere(brokers ...Agent) *Sphere {
	// declare agent
	var broker Agent
	// set declared agent if parameter exists
	for _, i := range brokers {
		broker = i
		break
	}
	if broker == nil {
		broker = NewSimpleBroker()
	}
	// creates sphere instance
	sphere := &Sphere{
		agent:       broker,
		connections: NewConnectionMap(),
		channels:    NewChannelMap(),
		models:      NewChannelModelMap(),
	}
	return sphere
}

// Sphere represents an entire Websocket instance
type Sphere struct {
	// a broker agent
	agent Agent
	// list of active connections
	connections ConnectionMap
	// list of channels
	channels ChannelMap
	// list of models
	models ChannelModelMap
}

// Handler handles and creates websocket connection
func (sphere *Sphere) Handler(w http.ResponseWriter, r *http.Request) {
	if conn, err := NewConnection(w, r); err == nil {
		sphere.connections.Set(conn.id, conn)
		go sphere.write(conn)
		sphere.read(conn)
		defer func() {
			for item := range conn.channels.Iter() {
				channel := item.Val
				sphere.unsubscribe(channel.namespace, channel.room, conn)
			}
			sphere.connections.Remove(conn.id)
		}()
	}
}

// ChannelModels channel models
func (sphere *Sphere) ChannelModels(models ...IChannels) {
	for _, model := range models {
		if !sphere.models.Has(model.Namespace()) {
			sphere.models.Set(model.Namespace(), model)
		} else {
			panic(fmt.Sprintf("model \"%s\" is already existed", model.Namespace()))
		}
	}
}

func (sphere *Sphere) write(conn *Connection) {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
	}()
	for {
		select {
		case msg, ok := <-conn.send:
			if !ok {
				conn.emit(websocket.CloseMessage, []byte{})
				return
			}
			if err := conn.emit(websocket.TextMessage, msg); err != nil {
				return
			}
		case <-ticker.C:
			if err := conn.emit(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (sphere *Sphere) read(conn *Connection) {
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			return
		}
		if msg != nil {
			go sphere.process(conn, msg)
		}
	}
}

func (sphere *Sphere) process(conn *Connection, msg []byte) {
	p, err := ParsePacket(msg)
	if err != nil {
		fmt.Printf("Error: %v - %v", err.Error(), string(msg[:]))
		return
	}
	if p != nil {
		switch p.Type {
		case PacketTypeChannel:
			if p.Namespace != "" && p.Room != "" {
				p.Machine = sphere.agent.ID()
				sphere.publish(p, conn)
			} else {
				p.Error = ErrBadScheme
				if json, err := p.toJSON(); err == nil {
					conn.send <- json
				}
			}
		case PacketTypeSubscribe:
			if p.Namespace != "" && p.Room != "" {
				err := sphere.subscribe(p.Namespace, p.Room, conn)
				r := p.Response()
				r.SetError(err)
				conn.emit(TextMessage, r, true)
			} else {
				p.Error = ErrBadScheme
				if json, err := p.toJSON(); err == nil {
					conn.send <- json
				}
			}
		case PacketTypeUnsubscribe:
			if p.Namespace != "" && p.Room != "" {
				if sphere.models.Has(p.Namespace) {
					sphere.unsubscribe(p.Namespace, p.Room, conn)
				} else {
					p.Error = ErrUnsupportedNamespace
					if json, err := p.toJSON(); err == nil {
						conn.send <- json
					}
				}
			} else {
				p.Error = ErrBadScheme
				if json, err := p.toJSON(); err == nil {
					conn.send <- json
				}
			}
		case PacketTypePing:
			r := p.Response()
			if json, err := r.toJSON(); err == nil {
				conn.send <- json
			}
		case PacketTypeMessage:
			conn.receive <- p
		}
	}
}

func (sphere *Sphere) channel(namespace string, room string, autoCreateOpts ...bool) *Channel {
	c := make(chan *Channel)
	autoCreateOpt := false
	for _, opt := range autoCreateOpts {
		autoCreateOpt = opt
		break
	}
	name := sphere.agent.ChannelName(namespace, room)
	go func() {
		if tmp, ok := sphere.channels.Get(name); ok {
			c <- tmp
		} else {
			if autoCreateOpt {
				channel := NewChannel(namespace, room)
				sphere.channels.Set(name, channel)
				c <- channel
			} else {
				c <- nil
			}
		}
	}()
	return <-c
}

func (sphere *Sphere) subscribe(namespace string, room string, conn *Connection) error {
	var model IChannels
	if !sphere.models.Has(namespace) {
		return ErrUnsupportedNamespace
	}
	if tmp, ok := sphere.models.Get(namespace); ok {
		model = tmp
	} else {
		return ErrUnsupportedNamespace
	}
	if accept := model.Subscribe(room, conn); !accept {
		return ErrNotAuthorized
	}
	channel := sphere.channel(namespace, room, true)
	if channel == nil {
		return ErrBadStatus
	}
	if err := channel.subscribe(conn); err == nil && !sphere.agent.IsSubscribed(channel.namespace, channel.room) {
		go sphere.agent.OnSubscribe(channel)
	}
	return nil
}

func (sphere *Sphere) unsubscribe(namespace string, room string, conn *Connection) error {
	var model IChannels
	if !sphere.models.Has(namespace) {
		return ErrUnsupportedNamespace
	}
	if tmp, ok := sphere.models.Get(namespace); ok {
		model = tmp
	} else {
		return ErrUnsupportedNamespace
	}
	if success := model.Disconnect(room, conn); !success {
		return ErrBadStatus
	}
	channel := sphere.channel(namespace, room, false)
	if channel == nil {
		return ErrBadStatus
	}
	if err := channel.unsubscribe(conn); err == nil && channel.connections.Count() == 0 {
		if sphere.agent.IsSubscribed(channel.namespace, channel.room) {
			go sphere.agent.OnUnsubscribe(channel)
		}
	}
	return nil
}

func (sphere *Sphere) publish(p *Packet, conn *Connection) error {
	var model IChannels
	if !sphere.models.Has(p.Namespace) {
		return ErrUnsupportedNamespace
	}
	if tmp, ok := sphere.models.Get(p.Namespace); ok {
		model = tmp
	} else {
		return ErrUnsupportedNamespace
	}
	channel := sphere.channel(p.Namespace, p.Room)
	if channel == nil {
		return ErrBadStatus
	}
	if !channel.isSubscribed(conn) {
		return ErrNotAuthorized
	}
	msg := p.Message
	if msg == nil || msg.Event == "" {
		return ErrBadScheme
	}
	res, err := model.Receive(msg.Event, msg.Data)
	if err != nil {
		return err
	}
	d := p.Response()
	if res != "" {
		d.Message.Data = res
	}
	if sphere.agent.IsSubscribed(channel.namespace, channel.room) {
		sphere.agent.OnPublish(channel, d)
	}
	return nil
}
<|endoftext|>