<commit_before>package aero

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/OneOfOne/xxhash"
	"github.com/aerogo/session"
	"github.com/fatih/color"
	"github.com/julienschmidt/httprouter"
	cache "github.com/patrickmn/go-cache"
	"github.com/tomasen/realip"
	"github.com/valyala/fasthttp"
)

// This should be close to the MTU size of a TCP packet.
// Regarding performance it makes no sense to compress smaller files.
// Bandwidth can be saved however the savings are minimal for small files
// and the overhead of compressing can lead up to a 75% reduction
// in server speed under high load. Therefore in this case
// we're trying to optimize for performance, not bandwidth.
const gzipThreshold = 1450

const (
	serverHeader                  = "Server"
	server                        = "Aero"
	cacheControlHeader            = "Cache-Control"
	cacheControlAlwaysValidate    = "must-revalidate"
	cacheControlMedia             = "public, max-age=864000"
	contentTypeOptionsHeader      = "X-Content-Type-Options"
	contentTypeOptions            = "nosniff"
	xssProtectionHeader           = "X-XSS-Protection"
	xssProtection                 = "1; mode=block"
	etagHeader                    = "ETag"
	contentTypeHeader             = "Content-Type"
	contentTypeHTML               = "text/html; charset=utf-8"
	contentTypeJavaScript         = "application/javascript; charset=utf-8"
	contentTypeJSON               = "application/json; charset=utf-8"
	contentTypePlainText          = "text/plain; charset=utf-8"
	contentEncodingHeader         = "Content-Encoding"
	contentEncodingGzip           = "gzip"
	contentLengthHeader           = "Content-Length"
	responseTimeHeader            = "X-Response-Time"
	ifNoneMatchHeader             = "If-None-Match"
	xFrameOptionsHeader           = "X-Frame-Options"
	xFrameOptions                 = "SAMEORIGIN"
	referrerPolicyHeader          = "Referrer-Policy"
	referrerPolicySameOrigin      = "no-referrer"
	strictTransportSecurityHeader = "Strict-Transport-Security"
	strictTransportSecurity       = "max-age=31536000; includeSubDomains; preload"
	contentSecurityPolicyHeader   = "Content-Security-Policy"
)

// Context ...
type Context struct {
	// net/http
	request  *http.Request
	response http.ResponseWriter
	params   httprouter.Params

	// A pointer to the application this request occured on.
	App *Application

	// Status code
	StatusCode int

	// Custom data
	Data interface{}

	// User session
	session *session.Session
}

// Request returns the HTTP request.
func (ctx *Context) Request() Request {
	return Request{
		inner: ctx.request,
	}
}

// Response returns the HTTP response.
func (ctx *Context) Response() Response {
	return Response{
		inner: ctx.response,
	}
}

// Session returns the session of the context or creates and caches a new session.
func (ctx *Context) Session() *session.Session {
	// Return cached session if available.
	if ctx.session != nil {
		return ctx.session
	}

	// Check if the client has a session cookie already.
	cookie, err := ctx.request.Cookie("sid")

	if err == nil {
		sid := cookie.Value

		if session.IsValidID(sid) {
			ctx.session, err = ctx.App.Sessions.Store.Get(sid)

			if err != nil {
				color.Red(err.Error())
			}

			if ctx.session != nil {
				return ctx.session
			}
		}
	}

	// Create a new session
	ctx.session = ctx.App.Sessions.New()

	// Create a session cookie in the client
	ctx.createSessionCookie()

	return ctx.session
}

// createSessionCookie creates a session cookie in the client.
func (ctx *Context) createSessionCookie() {
	sessionCookie := http.Cookie{
		Name:     "sid",
		Value:    ctx.session.ID(),
		HttpOnly: true,
		Secure:   true,
		MaxAge:   ctx.App.Sessions.Duration,
		Path:     "/",
	}

	http.SetCookie(ctx.response, &sessionCookie)

	// HACK: Add SameSite attribute
	// Remove this once it's available inside http.Cookie
	// cookieData := ctx.response.Header().Get("Set-Cookie")
	// cookieData += "; SameSite=lax"
	// ctx.response.Header().Set("Set-Cookie", cookieData)
}

// HasSession indicates whether the client has a valid session or not.
func (ctx *Context) HasSession() bool {
	if ctx.session != nil {
		return true
	}

	cookie, err := ctx.request.Cookie("sid")

	if err != nil || !session.IsValidID(cookie.Value) {
		return false
	}

	ctx.session, _ = ctx.App.Sessions.Store.Get(cookie.Value)

	return ctx.session != nil
}

// JSON encodes the object to a JSON string and responds.
func (ctx *Context) JSON(value interface{}) string {
	bytes, _ := json.Marshal(value)

	ctx.response.Header().Set(contentTypeHeader, contentTypeJSON)
	return string(bytes)
}

// HTML sends a HTML string.
func (ctx *Context) HTML(html string) string {
	ctx.response.Header().Set(contentTypeHeader, contentTypeHTML)
	ctx.response.Header().Set(contentTypeOptionsHeader, contentTypeOptions)
	ctx.response.Header().Set(xssProtectionHeader, xssProtection)
	// ctx.response.Header().Set(xFrameOptionsHeader, xFrameOptions)
	ctx.response.Header().Set(referrerPolicyHeader, referrerPolicySameOrigin)

	if ctx.App.Security.Certificate != "" {
		ctx.response.Header().Set(strictTransportSecurityHeader, strictTransportSecurity)
		ctx.response.Header().Set(contentSecurityPolicyHeader, ctx.App.contentSecurityPolicy)
	}

	return html
}

// Text sends a plain text string.
func (ctx *Context) Text(text string) string {
	ctx.response.Header().Set(contentTypeHeader, contentTypePlainText)
	return text
}

// JavaScript sends a script.
func (ctx *Context) JavaScript(code string) string {
	ctx.response.Header().Set(contentTypeHeader, contentTypeJavaScript)
	return code
}

// File sends the contents of a local file and determines its mime type by extension.
func (ctx *Context) File(file string) string {
	extension := filepath.Ext(file)
	mimeType := mime.TypeByExtension(extension)
	data, _ := ioutil.ReadFile(file)

	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}

	ctx.response.Header().Set(contentTypeHeader, mimeType)
	return string(data)
}

// TryWebP tries to serve a WebP image but will fall back to the specified extension if needed.
func (ctx *Context) TryWebP(path string, extension string) string {
	if ctx.CanUseWebP() {
		extension = ".webp"
	}

	return ctx.File(path + extension)
}

// Error should be used for sending error messages to the user.
func (ctx *Context) Error(statusCode int, explanation string, err error) string {
	ctx.StatusCode = statusCode
	ctx.response.Header().Set(contentTypeHeader, contentTypeHTML)

	if err != nil {
		detailed := err.Error()
		color.Red(detailed)
		return fmt.Sprintf("%s (%s)", explanation, detailed)
	}

	return explanation
}

// URI returns the relative path, e.g. /blog/post/123.
func (ctx *Context) URI() string {
	return ctx.request.URL.Path
}

// SetURI sets the relative path, e.g. /blog/post/123.
func (ctx *Context) SetURI(b string) {
	ctx.request.URL.Path = b
}

// Get retrieves an URL parameter.
func (ctx *Context) Get(param string) string {
	return ctx.params.ByName(param)
}

// GetInt retrieves an URL parameter as an integer.
func (ctx *Context) GetInt(param string) (int, error) {
	return strconv.Atoi(ctx.Get(param))
}

// RealIP tries to determine the real IP address of the request.
func (ctx *Context) RealIP() string {
	return realip.RealIP(ctx.request)
}

// UserAgent retrieves the user agent for the given request.
func (ctx *Context) UserAgent() string {
	ctx.request.URL.Query()
	return ctx.request.UserAgent()
}

// Query retrieves the value for the given URL query parameter.
func (ctx *Context) Query(param string) string {
	return ctx.request.URL.Query().Get(param)
}

// Redirect redirects to the given URL using status code 302.
func (ctx *Context) Redirect(url string) string {
	ctx.StatusCode = http.StatusFound
	ctx.response.Header().Set("Location", url)
	return ""
}

// RedirectPermanently redirects to the given URL and indicates that this is a permanent change using status code 301.
func (ctx *Context) RedirectPermanently(url string) string {
	ctx.StatusCode = http.StatusPermanentRedirect
	ctx.response.Header().Set("Location", url)
	return ""
}

// CanUseWebP checks the Accept header to find out if WebP is supported by the client's browser.
func (ctx *Context) CanUseWebP() bool {
	accept := ctx.request.Header.Get("Accept")

	if strings.Index(accept, "image/webp") != -1 {
		return true
	}

	return false
}

// IsMediaResponse returns whether the given context has already set its content type to a media type.
func (ctx *Context) IsMediaResponse() bool {
	contentType := ctx.response.Header().Get(contentTypeHeader)
	return strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "video/")
}

// respond responds either with raw code or gzipped if the
// code length is greater than the gzip threshold.
func (ctx *Context) respond(code string) {
	ctx.respondBytes(StringToBytesUnsafe(code))
}

// respondBytes responds either with raw code or gzipped if the
// code length is greater than the gzip threshold. Requires a byte slice.
func (ctx *Context) respondBytes(b []byte) {
	response := ctx.response
	header := response.Header()
	isMedia := ctx.IsMediaResponse()

	// Headers
	if isMedia {
		header.Set(cacheControlHeader, cacheControlMedia)
	} else {
		header.Set(cacheControlHeader, cacheControlAlwaysValidate)
		header.Set(serverHeader, server)
		// header.Set(responseTimeHeader, strconv.FormatInt(time.Since(ctx.start).Nanoseconds()/1000, 10)+" us")
	}

	// Small response
	if len(b) < gzipThreshold {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// ETag generation
	h := xxhash.NewS64(0)
	h.Write(b)
	etag := strconv.FormatUint(h.Sum64(), 16)

	// If client cache is up to date, send 304 with no response body.
	clientETag := ctx.request.Header.Get(ifNoneMatchHeader)

	if etag == clientETag {
		response.WriteHeader(304)
		return
	}

	// Set ETag
	header.Set(etagHeader, etag)

	// No GZip?
	if !ctx.App.Config.GZip || isMedia {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// GZip
	header.Set(contentEncodingHeader, contentEncodingGzip)

	if ctx.App.Config.GZipCache {
		cachedResponse, found := ctx.App.gzipCache.Get(etag)

		if found {
			cachedResponseBytes := cachedResponse.([]byte)
			header.Set(contentLengthHeader, strconv.Itoa(len(cachedResponseBytes)))
			response.WriteHeader(ctx.StatusCode)
			response.Write(cachedResponseBytes)
			return
		}
	}

	var buffer bytes.Buffer
	writer := bufio.NewWriter(&buffer)
	fasthttp.WriteGzipLevel(writer, b, 9)
	writer.Flush()
	gzippedBytes := buffer.Bytes()

	header.Set(contentLengthHeader, strconv.Itoa(len(gzippedBytes)))
	response.WriteHeader(ctx.StatusCode)
	response.Write(gzippedBytes)

	if ctx.App.Config.GZipCache {
		ctx.App.gzipCache.Set(etag, gzippedBytes, cache.DefaultExpiration)
	}
}
<commit_msg>Added error when JSON encoding fails<commit_after>package aero

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/OneOfOne/xxhash"
	"github.com/aerogo/session"
	"github.com/fatih/color"
	"github.com/julienschmidt/httprouter"
	cache "github.com/patrickmn/go-cache"
	"github.com/tomasen/realip"
	"github.com/valyala/fasthttp"
)

// This should be close to the MTU size of a TCP packet.
// Regarding performance it makes no sense to compress smaller files.
// Bandwidth can be saved however the savings are minimal for small files
// and the overhead of compressing can lead up to a 75% reduction
// in server speed under high load. Therefore in this case
// we're trying to optimize for performance, not bandwidth.
const gzipThreshold = 1450

const (
	serverHeader                  = "Server"
	server                        = "Aero"
	cacheControlHeader            = "Cache-Control"
	cacheControlAlwaysValidate    = "must-revalidate"
	cacheControlMedia             = "public, max-age=864000"
	contentTypeOptionsHeader      = "X-Content-Type-Options"
	contentTypeOptions            = "nosniff"
	xssProtectionHeader           = "X-XSS-Protection"
	xssProtection                 = "1; mode=block"
	etagHeader                    = "ETag"
	contentTypeHeader             = "Content-Type"
	contentTypeHTML               = "text/html; charset=utf-8"
	contentTypeJavaScript         = "application/javascript; charset=utf-8"
	contentTypeJSON               = "application/json; charset=utf-8"
	contentTypePlainText          = "text/plain; charset=utf-8"
	contentEncodingHeader         = "Content-Encoding"
	contentEncodingGzip           = "gzip"
	contentLengthHeader           = "Content-Length"
	responseTimeHeader            = "X-Response-Time"
	ifNoneMatchHeader             = "If-None-Match"
	xFrameOptionsHeader           = "X-Frame-Options"
	xFrameOptions                 = "SAMEORIGIN"
	referrerPolicyHeader          = "Referrer-Policy"
	referrerPolicySameOrigin      = "no-referrer"
	strictTransportSecurityHeader = "Strict-Transport-Security"
	strictTransportSecurity       = "max-age=31536000; includeSubDomains; preload"
	contentSecurityPolicyHeader   = "Content-Security-Policy"
)

// Context ...
type Context struct {
	// net/http
	request  *http.Request
	response http.ResponseWriter
	params   httprouter.Params

	// A pointer to the application this request occured on.
	App *Application

	// Status code
	StatusCode int

	// Custom data
	Data interface{}

	// User session
	session *session.Session
}

// Request returns the HTTP request.
func (ctx *Context) Request() Request {
	return Request{
		inner: ctx.request,
	}
}

// Response returns the HTTP response.
func (ctx *Context) Response() Response {
	return Response{
		inner: ctx.response,
	}
}

// Session returns the session of the context or creates and caches a new session.
func (ctx *Context) Session() *session.Session {
	// Return cached session if available.
	if ctx.session != nil {
		return ctx.session
	}

	// Check if the client has a session cookie already.
	cookie, err := ctx.request.Cookie("sid")

	if err == nil {
		sid := cookie.Value

		if session.IsValidID(sid) {
			ctx.session, err = ctx.App.Sessions.Store.Get(sid)

			if err != nil {
				color.Red(err.Error())
			}

			if ctx.session != nil {
				return ctx.session
			}
		}
	}

	// Create a new session
	ctx.session = ctx.App.Sessions.New()

	// Create a session cookie in the client
	ctx.createSessionCookie()

	return ctx.session
}

// createSessionCookie creates a session cookie in the client.
func (ctx *Context) createSessionCookie() {
	sessionCookie := http.Cookie{
		Name:     "sid",
		Value:    ctx.session.ID(),
		HttpOnly: true,
		Secure:   true,
		MaxAge:   ctx.App.Sessions.Duration,
		Path:     "/",
	}

	http.SetCookie(ctx.response, &sessionCookie)

	// HACK: Add SameSite attribute
	// Remove this once it's available inside http.Cookie
	// cookieData := ctx.response.Header().Get("Set-Cookie")
	// cookieData += "; SameSite=lax"
	// ctx.response.Header().Set("Set-Cookie", cookieData)
}

// HasSession indicates whether the client has a valid session or not.
func (ctx *Context) HasSession() bool {
	if ctx.session != nil {
		return true
	}

	cookie, err := ctx.request.Cookie("sid")

	if err != nil || !session.IsValidID(cookie.Value) {
		return false
	}

	ctx.session, _ = ctx.App.Sessions.Store.Get(cookie.Value)

	return ctx.session != nil
}

// JSON encodes the object to a JSON string and responds.
func (ctx *Context) JSON(value interface{}) string {
	ctx.response.Header().Set(contentTypeHeader, contentTypeJSON)

	bytes, err := json.Marshal(value)

	if err != nil {
		ctx.StatusCode = http.StatusInternalServerError
		return `{"error": "Could not encode object to JSON"}`
	}

	return string(bytes)
}

// HTML sends a HTML string.
func (ctx *Context) HTML(html string) string {
	ctx.response.Header().Set(contentTypeHeader, contentTypeHTML)
	ctx.response.Header().Set(contentTypeOptionsHeader, contentTypeOptions)
	ctx.response.Header().Set(xssProtectionHeader, xssProtection)
	// ctx.response.Header().Set(xFrameOptionsHeader, xFrameOptions)
	ctx.response.Header().Set(referrerPolicyHeader, referrerPolicySameOrigin)

	if ctx.App.Security.Certificate != "" {
		ctx.response.Header().Set(strictTransportSecurityHeader, strictTransportSecurity)
		ctx.response.Header().Set(contentSecurityPolicyHeader, ctx.App.contentSecurityPolicy)
	}

	return html
}

// Text sends a plain text string.
func (ctx *Context) Text(text string) string {
	ctx.response.Header().Set(contentTypeHeader, contentTypePlainText)
	return text
}

// JavaScript sends a script.
func (ctx *Context) JavaScript(code string) string {
	ctx.response.Header().Set(contentTypeHeader, contentTypeJavaScript)
	return code
}

// File sends the contents of a local file and determines its mime type by extension.
func (ctx *Context) File(file string) string {
	extension := filepath.Ext(file)
	mimeType := mime.TypeByExtension(extension)
	data, _ := ioutil.ReadFile(file)

	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}

	ctx.response.Header().Set(contentTypeHeader, mimeType)
	return string(data)
}

// TryWebP tries to serve a WebP image but will fall back to the specified extension if needed.
func (ctx *Context) TryWebP(path string, extension string) string {
	if ctx.CanUseWebP() {
		extension = ".webp"
	}

	return ctx.File(path + extension)
}

// Error should be used for sending error messages to the user.
func (ctx *Context) Error(statusCode int, explanation string, err error) string {
	ctx.StatusCode = statusCode
	ctx.response.Header().Set(contentTypeHeader, contentTypeHTML)

	if err != nil {
		detailed := err.Error()
		color.Red(detailed)
		return fmt.Sprintf("%s (%s)", explanation, detailed)
	}

	return explanation
}

// URI returns the relative path, e.g. /blog/post/123.
func (ctx *Context) URI() string {
	return ctx.request.URL.Path
}

// SetURI sets the relative path, e.g. /blog/post/123.
func (ctx *Context) SetURI(b string) {
	ctx.request.URL.Path = b
}

// Get retrieves an URL parameter.
func (ctx *Context) Get(param string) string {
	return ctx.params.ByName(param)
}

// GetInt retrieves an URL parameter as an integer.
func (ctx *Context) GetInt(param string) (int, error) {
	return strconv.Atoi(ctx.Get(param))
}

// RealIP tries to determine the real IP address of the request.
func (ctx *Context) RealIP() string {
	return realip.RealIP(ctx.request)
}

// UserAgent retrieves the user agent for the given request.
func (ctx *Context) UserAgent() string {
	ctx.request.URL.Query()
	return ctx.request.UserAgent()
}

// Query retrieves the value for the given URL query parameter.
func (ctx *Context) Query(param string) string {
	return ctx.request.URL.Query().Get(param)
}

// Redirect redirects to the given URL using status code 302.
func (ctx *Context) Redirect(url string) string {
	ctx.StatusCode = http.StatusFound
	ctx.response.Header().Set("Location", url)
	return ""
}

// RedirectPermanently redirects to the given URL and indicates that this is a permanent change using status code 301.
func (ctx *Context) RedirectPermanently(url string) string {
	ctx.StatusCode = http.StatusPermanentRedirect
	ctx.response.Header().Set("Location", url)
	return ""
}

// CanUseWebP checks the Accept header to find out if WebP is supported by the client's browser.
func (ctx *Context) CanUseWebP() bool {
	accept := ctx.request.Header.Get("Accept")

	if strings.Index(accept, "image/webp") != -1 {
		return true
	}

	return false
}

// IsMediaResponse returns whether the given context has already set its content type to a media type.
func (ctx *Context) IsMediaResponse() bool {
	contentType := ctx.response.Header().Get(contentTypeHeader)
	return strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "video/")
}

// respond responds either with raw code or gzipped if the
// code length is greater than the gzip threshold.
func (ctx *Context) respond(code string) {
	ctx.respondBytes(StringToBytesUnsafe(code))
}

// respondBytes responds either with raw code or gzipped if the
// code length is greater than the gzip threshold. Requires a byte slice.
func (ctx *Context) respondBytes(b []byte) {
	response := ctx.response
	header := response.Header()
	isMedia := ctx.IsMediaResponse()

	// Headers
	if isMedia {
		header.Set(cacheControlHeader, cacheControlMedia)
	} else {
		header.Set(cacheControlHeader, cacheControlAlwaysValidate)
		header.Set(serverHeader, server)
		// header.Set(responseTimeHeader, strconv.FormatInt(time.Since(ctx.start).Nanoseconds()/1000, 10)+" us")
	}

	// Small response
	if len(b) < gzipThreshold {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// ETag generation
	h := xxhash.NewS64(0)
	h.Write(b)
	etag := strconv.FormatUint(h.Sum64(), 16)

	// If client cache is up to date, send 304 with no response body.
	clientETag := ctx.request.Header.Get(ifNoneMatchHeader)

	if etag == clientETag {
		response.WriteHeader(304)
		return
	}

	// Set ETag
	header.Set(etagHeader, etag)

	// No GZip?
	if !ctx.App.Config.GZip || isMedia {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// GZip
	header.Set(contentEncodingHeader, contentEncodingGzip)

	if ctx.App.Config.GZipCache {
		cachedResponse, found := ctx.App.gzipCache.Get(etag)

		if found {
			cachedResponseBytes := cachedResponse.([]byte)
			header.Set(contentLengthHeader, strconv.Itoa(len(cachedResponseBytes)))
			response.WriteHeader(ctx.StatusCode)
			response.Write(cachedResponseBytes)
			return
		}
	}

	var buffer bytes.Buffer
	writer := bufio.NewWriter(&buffer)
	fasthttp.WriteGzipLevel(writer, b, 9)
	writer.Flush()
	gzippedBytes := buffer.Bytes()

	header.Set(contentLengthHeader, strconv.Itoa(len(gzippedBytes)))
	response.WriteHeader(ctx.StatusCode)
	response.Write(gzippedBytes)

	if ctx.App.Config.GZipCache {
		ctx.App.gzipCache.Set(etag, gzippedBytes, cache.DefaultExpiration)
	}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/Sirupsen/logrus"
	"github.com/citadel/citadel"
	"github.com/citadel/citadel/repository"
	"github.com/citadel/citadel/utils"
	"github.com/codegangsta/cli"
	"github.com/samalba/dockerclient"
)

type (
	HostEngine struct {
		client     *dockerclient.DockerClient
		repository *repository.Repository
		id         string
		listenAddr string
	}
)

var hostCommand = cli.Command{
	Name:   "host",
	Usage:  "run the host and connect it to the cluster",
	Action: hostAction,
	Flags: []cli.Flag{
		cli.StringFlag{"host-id", "", "specify host id (default: detected)"},
		cli.StringFlag{"region", "", "region where the host is running"},
		cli.StringFlag{"addr", "", "external ip address for the host"},
		cli.StringFlag{"docker", "unix:///var/run/docker.sock", "docker remote ip address"},
		cli.IntFlag{"cpus", -1, "number of cpus available to the host"},
		cli.IntFlag{"memory", -1, "number of mb of memory available to the host"},
		cli.StringFlag{"listen, l", ":8787", "listen address"},
	},
}

func hostAction(context *cli.Context) {
	var (
		cpus       = context.Int("cpus")
		memory     = context.Int("memory")
		addr       = context.String("addr")
		region     = context.String("region")
		hostId     = context.String("host-id")
		listenAddr = context.String("listen")
	)
	if hostId == "" {
		id, err := utils.GetMachineID()
		if err != nil {
			logger.WithField("error", err).Fatal("unable to read machine id")
		}
		hostId = id
	}

	switch {
	case cpus < 1:
		logger.Fatal("cpus must have a value")
	case memory < 1:
		logger.Fatal("memory must have a value")
	case addr == "":
		logger.Fatal("addr must have a value")
	case region == "":
		logger.Fatal("region must have a value")
	}

	machines := strings.Split(context.GlobalString("etcd-machines"), ",")
	r := repository.New(machines, "citadel")

	host := &citadel.Host{
		ID:     hostId,
		Memory: memory,
		Cpus:   cpus,
		Addr:   addr,
		Region: region,
	}

	if err := r.SaveHost(host); err != nil {
		logger.WithField("error", err).Fatal("unable to save host")
	}
	defer r.DeleteHost(hostId)

	client, err := dockerclient.NewDockerClient(context.String("docker"))
	if err != nil {
		logger.WithField("error", err).Fatal("unable to connect to docker")
	}

	hostEngine := &HostEngine{
		client:     client,
		repository: r,
		id:         hostId,
		listenAddr: listenAddr,
	}
	// start
	go hostEngine.run()
	// watch for operations
	go hostEngine.watch()
	// handle stop signal
	hostEngine.waitForInterrupt()
}

func (eng *HostEngine) waitForInterrupt() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
	for _ = range sigChan {
		// stop engine
		eng.stop()
		os.Exit(0)
	}
}

func (eng *HostEngine) run() {
	logger.Info("Starting Citadel")
	if err := eng.loadContainers(); err != nil {
		logger.WithField("error", err).Fatal("unable to load containers")
	}

	// listen for events
	eng.client.StartMonitorEvents(eng.dockerEventHandler)

	if err := http.ListenAndServe(eng.listenAddr, nil); err != nil {
		logger.WithField("error", err).Fatal("unable to listen on http")
	}
}

func (eng *HostEngine) stop() {
	logger.Info("Stopping")
	// remove host from repository
	eng.repository.DeleteHost(eng.id)
}

func (eng *HostEngine) loadContainers() error {
	eng.repository.DeleteHostContainers(eng.id)

	containers, err := eng.client.ListContainers(true)
	if err != nil {
		return err
	}

	for _, c := range containers {
		cc, err := eng.generateContainerInfo(c)
		if err != nil {
			return err
		}
		if err := eng.repository.SaveContainer(cc); err != nil {
			return err
		}
	}

	return nil
}

func (eng *HostEngine) generateContainerInfo(cnt interface{}) (*citadel.Container, error) {
	c := cnt.(dockerclient.Container)
	info, err := eng.client.InspectContainer(c.Id)
	if err != nil {
		return nil, err
	}
	cc := &citadel.Container{
		ID:     info.Id,
		Image:  utils.CleanImageName(c.Image),
		HostID: eng.id,
		Cpus:   info.Config.CpuShares, // FIXME: not the right place, this is cpuset
	}

	if info.Config.Memory > 0 {
		cc.Memory = info.Config.Memory / 1024 / 1024
	}

	if info.State.Running {
		cc.State.Status = citadel.Running
	} else {
		cc.State.Status = citadel.Stopped
	}
	cc.State.ExitCode = info.State.ExitCode
	return cc, nil
}

func (eng *HostEngine) dockerEventHandler(event *dockerclient.Event, args ...interface{}) {
	switch event.Status {
	case "destroy":
		// remove container from repository
		if err := eng.repository.DeleteContainer(eng.id, event.Id); err != nil {
			logger.Warnf("Unable to remove container from repository: %s", err)
		}
	default:
		// reload containers into repository
		// when adding a single container, the Container struct is not
		// returned but instead ContainerInfo.  to keep the same
		// generateContainerInfo for a citadel container, i simply
		// re-run the loadContainers.  this can probably be improved.
		eng.loadContainers()
	}
}

func (eng *HostEngine) watch() {
	tickerChan := time.NewTicker(time.Millisecond * 2000).C
	for _ = range tickerChan {
		tasks, err := eng.repository.FetchTasks()
		if err != nil {
			logger.Fatal("unable to fetch queue: %s", err)
		}

		for _, task := range tasks {
			// filter this hosts tasks
			if task.Host == eng.id {
				go eng.taskHandler(task)
			}
		}
	}
}

func (eng *HostEngine) taskHandler(task *citadel.Task) {
	switch task.Command {
	case "run":
		logger.WithFields(logrus.Fields{
			"host": task.Host,
		}).Info("processing run task")

		eng.runHandler(task)
	case "restart":
		logger.WithFields(logrus.Fields{
			"host": task.Host,
		}).Info("processing restart task")

		eng.restartHandler(task)
	case "stop":
		logger.WithFields(logrus.Fields{
			"host": task.Host,
		}).Info("processing stop task")

		eng.stopHandler(task)
	case "destroy":
		logger.WithFields(logrus.Fields{
			"host": task.Host,
		}).Info("processing destroy task")

		eng.destroyHandler(task)
	default:
		logger.WithFields(logrus.Fields{
			"command": task.Command,
		}).Error("unknown task command")
	}
}

func (eng *HostEngine) runHandler(task *citadel.Task) {
	logger.WithFields(logrus.Fields{
		"host":      task.Host,
		"image":     task.Image,
		"cpus":      task.Cpus,
		"memory":    task.Memory,
		"instances": task.Instances,
	}).Info("running container")

	eng.repository.DeleteTask(task.ID)

	for i := 0; i < task.Instances; i++ {
		containerConfig := &dockerclient.ContainerConfig{
			Image:     task.Image,
			Memory:    task.Memory * 1024 * 1024,
			CpuShares: task.Cpus,
		}

		containerId, err := eng.client.CreateContainer(containerConfig, "")
		if err != nil {
			logger.WithFields(logrus.Fields{
				"err": err,
			}).Error("error creating container")
			return
		}

		if err := eng.client.StartContainer(containerId, nil); err != nil {
			logger.WithFields(logrus.Fields{
				"err": err,
			}).Error("error starting container")
			return
		}

		logger.WithFields(logrus.Fields{
			"host":  task.Host,
			"id":    containerId,
			"image": task.Image,
		}).Info("started container")
	}
}

func (eng *HostEngine) stopHandler(task *citadel.Task) {
	logger.WithFields(logrus.Fields{
		"host": task.Host,
		"id":   task.ContainerID,
	}).Info("stopping container")

	defer eng.repository.DeleteTask(task.ID)

	containerId := task.ContainerID
	if err := eng.client.StopContainer(containerId, 10); err != nil {
		logger.WithFields(logrus.Fields{
			"id":  containerId,
			"err": err,
		}).Error("error stopping container")
	}
}

func (eng *HostEngine) restartHandler(task *citadel.Task) {
	logger.WithFields(logrus.Fields{
		"host": task.Host,
		"id":   task.ContainerID,
	}).Info("restarting container")

	defer eng.repository.DeleteTask(task.ID)

	containerId := task.ContainerID
	if err := eng.client.RestartContainer(containerId, 10); err != nil {
		logger.WithFields(logrus.Fields{
			"containerId": containerId,
			"err":         err,
		}).Error("error restarting container")
	}
}

func (eng *HostEngine) destroyHandler(task *citadel.Task) {
	logger.WithFields(logrus.Fields{
		"host": task.Host,
		"id":   task.ContainerID,
	}).Info("destroying container")

	defer eng.repository.DeleteTask(task.ID)

	containerId := task.ContainerID
	if err := eng.client.KillContainer(containerId); err != nil {
		logger.WithFields(logrus.Fields{
			"containerId": containerId,
			"err":         err,
		}).Error("error killing container")
		return
	}

	if err := eng.client.RemoveContainer(containerId); err != nil {
		logger.WithFields(logrus.Fields{
			"containerId": containerId,
			"err":         err,
		}).Error("error removing container")
	}
}
<commit_msg>Remove duplicate logs<commit_after>package main

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

	"github.com/Sirupsen/logrus"
	"github.com/citadel/citadel"
	"github.com/citadel/citadel/repository"
	"github.com/citadel/citadel/utils"
	"github.com/codegangsta/cli"
	"github.com/samalba/dockerclient"
)

type (
	HostEngine struct {
		client     *dockerclient.DockerClient
		repository *repository.Repository
		id         string
		listenAddr string
	}
)

var hostCommand = cli.Command{
	Name:   "host",
	Usage:  "run the host and connect it to the cluster",
	Action: hostAction,
	Flags: []cli.Flag{
		cli.StringFlag{"host-id", "", "specify host id (default: detected)"},
		cli.StringFlag{"region", "", "region where the host is running"},
		cli.StringFlag{"addr", "", "external ip address for the host"},
		cli.StringFlag{"docker", "unix:///var/run/docker.sock", "docker remote ip address"},
		cli.IntFlag{"cpus", -1, "number of cpus available to the host"},
		cli.IntFlag{"memory", -1, "number of mb of memory available to the host"},
		cli.StringFlag{"listen, l", ":8787", "listen address"},
	},
}

func hostAction(context *cli.Context) {
	var (
		cpus       = context.Int("cpus")
		memory     = context.Int("memory")
		addr       = context.String("addr")
		region     = context.String("region")
		hostId     = context.String("host-id")
		listenAddr = context.String("listen")
	)

	if hostId == "" {
		id, err := utils.GetMachineID()
		if err != nil {
			logger.WithField("error", err).Fatal("unable to read machine id")
		}
		hostId = id
	}

	switch {
	case cpus < 1:
		logger.Fatal("cpus must have a value")
	case memory < 1:
		logger.Fatal("memory must have a value")
	case addr == "":
		logger.Fatal("addr must have a value")
	case region == "":
		logger.Fatal("region must have a value")
	}

	machines := strings.Split(context.GlobalString("etcd-machines"), ",")
	r := repository.New(machines, "citadel")

	host := &citadel.Host{
		ID:     hostId,
		Memory: memory,
		Cpus:   cpus,
		Addr:   addr,
		Region: region,
	}

	if err := r.SaveHost(host); err != nil {
		logger.WithField("error", err).Fatal("unable to save host")
	}
	defer r.DeleteHost(hostId)

	client, err := dockerclient.NewDockerClient(context.String("docker"))
	if err != nil {
		logger.WithField("error", err).Fatal("unable to connect to docker")
	}

	hostEngine := &HostEngine{
		client:     client,
		repository: r,
		id:         hostId,
		listenAddr: listenAddr,
	}

	go hostEngine.run()
	go hostEngine.watch()

	hostEngine.waitForInterrupt()
}

func (eng *HostEngine) waitForInterrupt() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
	for _ = range sigChan {
		// stop engine
		eng.stop()
		os.Exit(0)
	}
}

func (eng *HostEngine) run() {
	logger.Info("Starting Citadel")

	if err := eng.loadContainers(); err != nil {
		logger.WithField("error", err).Fatal("unable to load containers")
	}

	// listen for events
	eng.client.StartMonitorEvents(eng.dockerEventHandler)

	if err := http.ListenAndServe(eng.listenAddr, nil); err != nil {
		logger.WithField("error", err).Fatal("unable to listen on http")
	}
}

func (eng *HostEngine) stop() {
	logger.Info("Stopping")
	// remove host from repository
	eng.repository.DeleteHost(eng.id)
}

func (eng *HostEngine) loadContainers() error {
	eng.repository.DeleteHostContainers(eng.id)

	containers, err := eng.client.ListContainers(true)
	if err != nil {
		return err
	}

	for _, c := range containers {
		cc, err := eng.generateContainerInfo(c)
		if err != nil {
			return err
		}
		if err := eng.repository.SaveContainer(cc); err != nil {
			return err
		}
	}

	return nil
}

func (eng *HostEngine) generateContainerInfo(cnt interface{}) (*citadel.Container, error) {
	c := cnt.(dockerclient.Container)
	info, err := eng.client.InspectContainer(c.Id)
	if err != nil {
		return nil, err
	}
	cc := &citadel.Container{
		ID:     info.Id,
		Image:  utils.CleanImageName(c.Image),
		HostID: eng.id,
		Cpus:   info.Config.CpuShares, // FIXME: not the right place, this is cpuset
	}

	if info.Config.Memory > 0 {
		cc.Memory = info.Config.Memory / 1024 / 1024
	}

	if info.State.Running {
		cc.State.Status = citadel.Running
	} else {
		cc.State.Status = citadel.Stopped
	}
	cc.State.ExitCode = info.State.ExitCode
	return cc, nil
}

func (eng *HostEngine) dockerEventHandler(event *dockerclient.Event, args ...interface{}) {
	switch event.Status {
	case "destroy":
		// remove container from repository
		if err := eng.repository.DeleteContainer(eng.id, event.Id); err != nil {
			logger.Warnf("Unable to remove container from repository: %s", err)
		}
	default:
		// reload containers into repository
		// when adding a single container, the Container struct is not
		// returned but instead ContainerInfo.  to keep the same
		// generateContainerInfo for a citadel container, i simply
		// re-run the loadContainers.  this can probably be improved.
		eng.loadContainers()
	}
}

func (eng *HostEngine) watch() {
	tickerChan := time.NewTicker(time.Millisecond * 2000).C
	for _ = range tickerChan {
		tasks, err := eng.repository.FetchTasks()
		if err != nil {
			logger.Fatal("unable to fetch queue: %s", err)
		}

		for _, task := range tasks {
			// filter this hosts tasks
			if task.Host == eng.id {
				go eng.taskHandler(task)
			}
		}
	}
}

func (eng *HostEngine) taskHandler(task *citadel.Task) {
	switch task.Command {
	case "run":
		logger.WithFields(logrus.Fields{
			"host": task.Host,
		}).Info("processing run task")

		eng.runHandler(task)
	case "restart":
		logger.WithFields(logrus.Fields{
			"host": task.Host,
		}).Info("processing restart task")

		eng.restartHandler(task)
	case "stop":
		logger.WithFields(logrus.Fields{
			"host": task.Host,
		}).Info("processing stop task")

		eng.stopHandler(task)
	case "destroy":
		logger.WithFields(logrus.Fields{
			"host": task.Host,
		}).Info("processing destroy task")

		eng.destroyHandler(task)
	default:
		logger.WithFields(logrus.Fields{
			"command": task.Command,
		}).Error("unknown task command")
	}
}

func (eng *HostEngine) runHandler(task *citadel.Task) {
	eng.repository.DeleteTask(task.ID)

	for i := 0; i < task.Instances; i++ {
		containerConfig := &dockerclient.ContainerConfig{
			Image:     task.Image,
			Memory:    task.Memory * 1024 * 1024,
			CpuShares: task.Cpus,
		}

		containerId, err := eng.client.CreateContainer(containerConfig, "")
		if err != nil {
			logger.WithFields(logrus.Fields{
				"err": err,
			}).Error("error creating container")
			return
		}

		if err := eng.client.StartContainer(containerId, nil); err != nil {
			logger.WithFields(logrus.Fields{
				"err": err,
			}).Error("error starting container")
			return
		}

		logger.WithFields(logrus.Fields{
			"host":  task.Host,
			"id":    containerId,
			"image": task.Image,
		}).Info("started container")
	}
}

func (eng *HostEngine) stopHandler(task *citadel.Task) {
	defer eng.repository.DeleteTask(task.ID)

	containerId := task.ContainerID
	if err := eng.client.StopContainer(containerId, 10); err != nil {
		logger.WithFields(logrus.Fields{
			"id":  containerId,
			"err": err,
		}).Error("error stopping container")
	}
}

func (eng *HostEngine) restartHandler(task *citadel.Task) {
	defer eng.repository.DeleteTask(task.ID)

	containerId := task.ContainerID
	if err := eng.client.RestartContainer(containerId, 10); err != nil {
		logger.WithFields(logrus.Fields{
			"containerId": containerId,
			"err":         err,
		}).Error("error restarting container")
	}
}

func (eng *HostEngine) destroyHandler(task *citadel.Task) {
	defer eng.repository.DeleteTask(task.ID)

	containerId := task.ContainerID
	if err := eng.client.KillContainer(containerId); err != nil {
		logger.WithFields(logrus.Fields{
			"containerId": containerId,
			"err":         err,
		}).Error("error killing container")
		return
	}

	if err := eng.client.RemoveContainer(containerId); err != nil {
		logger.WithFields(logrus.Fields{
			"containerId": containerId,
			"err":         err,
		}).Error("error removing container")
	}
}
<|endoftext|>


<commit_before>package lmdb

import (
	"errors"
	mdb "github.com/szferi/gomdb"
)

// Thread Safety
// 1) NOTLS mode is used exclusively, which allows read txns to freely migrate across
//    threads and for a single thread to maintain multiple read txns. This enables mostly
//    care-free use of read txns.
// 2) Most objects can be safely called by a single caller from a single thread, and usually it
//    only makes sense to have a single caller, except in the case of Database.
// 3) Most Database methods are thread-safe, and may be called concurrently, except for
//    Database.Close().
// 4) A write txn may only be used from the thread it was created on.
// 5) A read-only txn can move across threads, but it cannot be used concurrently from multiple
//    threads.
// 6) Iterator is not thread-safe, but it does not make sense to use it on any thread except the
//    thread that currently owns its associated txn.
//
//-------------------------------------------------------------------------------------------------
//
// Best practice:
// 1) Use iterators only in the txn that they are created
// 2) DO NOT modify the memory slice from GetNoCopy
// 3) Make sure all read/write txns are finished before Database.Close().

const (
	// There is no penalty for making this huge.
	// If you are on a 32-bit system, use Open2 and specify a smaller map size.
	MAP_SIZE_DEFAULT uint64 = 64 * 1024 * 1024 * 1024 * 1024 // 64TB
)

type TransactionalRWer interface {
	TransactionalRW(func(*ReadWriteTxn) error) error
}

type Database struct {
	env *mdb.Env
	// In this package, a DBI is obtained only through Open/Open2, and is never closed until
	// Context.Close(), in which all dbis are closed automatically.
	buckets map[string]mdb.DBI
}

type Stat mdb.Stat
type Info mdb.Info

//--------------------------------- Database ------------------------------------------------------

func Version() string {
	return mdb.Version()
}

func Open(path string, buckets []string) (*Database, error) {
	return Open2(path, buckets, MAP_SIZE_DEFAULT)
}

func Open2(path string, buckets []string, maxMapSize uint64) (db *Database, err error) {
	// TODO (Potential bug):
	// From mdb_env_open's doc,
	//  "If this function fails, #mdb_env_close() must be called to discard the #MDB_env handle."
	// But mdb.NewEnv doesnot call mdb_env_close() when it fails, AND it just return nil as env.
	// Patch gomdb if this turns out to be a big issue.
	env, err := mdb.NewEnv()
	defer func() {
		if err != nil && env != nil {
			env.Close()
			db.env = nil
		}
	}()
	if err != nil {
		return
	}

	err = env.SetMapSize(maxMapSize)
	if err != nil {
		return
	}

	// http://www.openldap.org/lists/openldap-technical/201305/msg00176.html
	err = env.SetMaxDBs((mdb.DBI)(len(buckets)))
	if err != nil {
		return
	}

	MDB_NOTLS := uint(0x200000)
	MDB_NORDAHEAD := uint(0x800000)
	err = env.Open(path, MDB_NOTLS|MDB_NORDAHEAD, 0664)
	if err != nil {
		return
	}

	bucketCache := make(map[string]mdb.DBI)
	db = &Database{env, nil}

	err = db.TransactionalRW(func(txn *ReadWriteTxn) error {
		for _, name := range buckets {
			if name == "" {
				return errors.New("Bucket name is empty")
			}
			dbi, err := txn.txn.DBIOpen(&name, mdb.CREATE)
			if err != nil {
				return err
			} else {
				bucketCache[name] = dbi
			}
		}
		return nil
	})

	if err != nil {
		return
	} else {
		db.buckets = bucketCache
	}

	return
}

func (db *Database) Close() {
	if db.env != nil {
		db.env.Close() // all opened dbis are closed during this process
	}
}

func (db *Database) Stat() *Stat {
	stat, err := db.env.Stat()
	if err != nil { // Possible errors: EINVAL
		panic(err)
	}
	return (*Stat)(stat)
}

func (db *Database) Info() *Info {
	info, err := db.env.Info()
	if err != nil { // error when env == nil, so panic
		panic(err)
	}
	return (*Info)(info)
}

func (db *Database) TransactionalR(f func(ReadTxner)) {
	txn, err := db.env.BeginTxn(nil, mdb.RDONLY)
	if err != nil { // Possible Errors: MDB_PANIC, MDB_MAP_RESIZED, MDB_READERS_FULL, ENOMEM
		panic(err)
	}

	var panicF interface{} // panic from f
	rdTxn := ReadTxn{db.buckets, txn, nil}

	defer func() {
		for _, itr := range rdTxn.itrs {
			itr.Close() // no panic
		}
		rdTxn.itrs = nil

		txn.Abort()
		if panicF != nil {
			panic(panicF) // re-panic
		}
	}()

	func() {
		defer func() {
			panicF = recover()
		}()
		f(&rdTxn)
	}()
}

func (db *Database) TransactionalRW(f func(*ReadWriteTxn) error) (err error) {
	txn, err := db.env.BeginTxn(nil, 0)
	if err != nil { // Possible Errors: MDB_PANIC, MDB_MAP_RESIZED, MDB_READERS_FULL, ENOMEM
		panic(err)
	}

	var panicF interface{} // panic from f
	rwCtx := ReadWriteTxn{db.env, &ReadTxn{db.buckets, txn, nil}}

	defer func() {
		for _, itr := range rwCtx.itrs {
			itr.Close() // no panic
		}
		rwCtx.itrs = nil

		if err == nil && panicF == nil {
			e := txn.Commit()
			if e != nil { // Possible errors: EINVAL, ENOSPEC, EIO, ENOMEM
				panic(e)
			}
		} else {
			txn.Abort()
			if panicF != nil {
				panic(panicF) // re-panic
			}
		}
	}()

	func() {
		defer func() {
			panicF = recover()
		}()
		err = f(&rwCtx)
	}()

	return
}
<commit_msg>rename TransactionalRWer to RWTxnCreator<commit_after>package lmdb

import (
	"errors"
	mdb "github.com/szferi/gomdb"
)

// Thread Safety
// 1) NOTLS mode is used exclusively, which allows read txns to freely migrate across
//    threads and for a single thread to maintain multiple read txns. This enables mostly
//    care-free use of read txns.
// 2) Most objects can be safely called by a single caller from a single thread, and usually it
//    only makes sense to have a single caller, except in the case of Database.
// 3) Most Database methods are thread-safe, and may be called concurrently, except for
//    Database.Close().
// 4) A write txn may only be used from the thread it was created on.
// 5) A read-only txn can move across threads, but it cannot be used concurrently from multiple
//    threads.
// 6) Iterator is not thread-safe, but it does not make sense to use it on any thread except the
//    thread that currently owns its associated txn.
//
//-------------------------------------------------------------------------------------------------
//
// Best practice:
// 1) Use iterators only in the txn that they are created
// 2) DO NOT modify the memory slice from GetNoCopy
// 3) Make sure all read/write txns are finished before Database.Close().

const (
	// There is no penalty for making this huge.
	// If you are on a 32-bit system, use Open2 and specify a smaller map size.
	MAP_SIZE_DEFAULT uint64 = 64 * 1024 * 1024 * 1024 * 1024 // 64TB
)

type RWTxnCreator interface {
	TransactionalRW(func(*ReadWriteTxn) error) error
}

type Database struct {
	env *mdb.Env
	// In this package, a DBI is obtained only through Open/Open2, and is never closed until
	// Context.Close(), in which all dbis are closed automatically.
	buckets map[string]mdb.DBI
}

type Stat mdb.Stat
type Info mdb.Info

//--------------------------------- Database ------------------------------------------------------

func Version() string {
	return mdb.Version()
}

func Open(path string, buckets []string) (*Database, error) {
	return Open2(path, buckets, MAP_SIZE_DEFAULT)
}

func Open2(path string, buckets []string, maxMapSize uint64) (db *Database, err error) {
	// TODO (Potential bug):
	// From mdb_env_open's doc,
	//  "If this function fails, #mdb_env_close() must be called to discard the #MDB_env handle."
	// But mdb.NewEnv doesnot call mdb_env_close() when it fails, AND it just return nil as env.
	// Patch gomdb if this turns out to be a big issue.
	env, err := mdb.NewEnv()
	defer func() {
		if err != nil && env != nil {
			env.Close()
			db.env = nil
		}
	}()
	if err != nil {
		return
	}

	err = env.SetMapSize(maxMapSize)
	if err != nil {
		return
	}

	// http://www.openldap.org/lists/openldap-technical/201305/msg00176.html
	err = env.SetMaxDBs((mdb.DBI)(len(buckets)))
	if err != nil {
		return
	}

	MDB_NOTLS := uint(0x200000)
	MDB_NORDAHEAD := uint(0x800000)
	err = env.Open(path, MDB_NOTLS|MDB_NORDAHEAD, 0664)
	if err != nil {
		return
	}

	bucketCache := make(map[string]mdb.DBI)
	db = &Database{env, nil}

	err = db.TransactionalRW(func(txn *ReadWriteTxn) error {
		for _, name := range buckets {
			if name == "" {
				return errors.New("Bucket name is empty")
			}
			dbi, err := txn.txn.DBIOpen(&name, mdb.CREATE)
			if err != nil {
				return err
			} else {
				bucketCache[name] = dbi
			}
		}
		return nil
	})

	if err != nil {
		return
	} else {
		db.buckets = bucketCache
	}

	return
}

func (db *Database) Close() {
	if db.env != nil {
		db.env.Close() // all opened dbis are closed during this process
	}
}

func (db *Database) Stat() *Stat {
	stat, err := db.env.Stat()
	if err != nil { // Possible errors: EINVAL
		panic(err)
	}
	return (*Stat)(stat)
}

func (db *Database) Info() *Info {
	info, err := db.env.Info()
	if err != nil { // error when env == nil, so panic
		panic(err)
	}
	return (*Info)(info)
}

func (db *Database) TransactionalR(f func(ReadTxner)) {
	txn, err := db.env.BeginTxn(nil, mdb.RDONLY)
	if err != nil { // Possible Errors: MDB_PANIC, MDB_MAP_RESIZED, MDB_READERS_FULL, ENOMEM
		panic(err)
	}

	var panicF interface{} // panic from f
	rdTxn := ReadTxn{db.buckets, txn, nil}

	defer func() {
		for _, itr := range rdTxn.itrs {
			itr.Close() // no panic
		}
		rdTxn.itrs = nil

		txn.Abort()
		if panicF != nil {
			panic(panicF) // re-panic
		}
	}()

	func() {
		defer func() {
			panicF = recover()
		}()
		f(&rdTxn)
	}()
}

func (db *Database) TransactionalRW(f func(*ReadWriteTxn) error) (err error) {
	txn, err := db.env.BeginTxn(nil, 0)
	if err != nil { // Possible Errors: MDB_PANIC, MDB_MAP_RESIZED, MDB_READERS_FULL, ENOMEM
		panic(err)
	}

	var panicF interface{} // panic from f
	rwCtx := ReadWriteTxn{db.env, &ReadTxn{db.buckets, txn, nil}}

	defer func() {
		for _, itr := range rwCtx.itrs {
			itr.Close() // no panic
		}
		rwCtx.itrs = nil

		if err == nil && panicF == nil {
			e := txn.Commit()
			if e != nil { // Possible errors: EINVAL, ENOSPEC, EIO, ENOMEM
				panic(e)
			}
		} else {
			txn.Abort()
			if panicF != nil {
				panic(panicF) // re-panic
			}
		}
	}()

	func() {
		defer func() {
			panicF = recover()
		}()
		err = f(&rwCtx)
	}()

	return
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha1"
	"database/sql"
	"fmt"
	"io"
	"log"
	"regexp"
	"strconv"

	_ "github.com/go-sql-driver/mysql"
)

func colPos(slice []string, value string) int {
	for p, v := range slice {
		if v == value {
			return p
		}
	}

	return -1
}

func getColByName(name string, cols []string, vals []interface{}) *string {
	if cmdi := colPos(cols, name); cmdi != -1 {
		if bytes, ok := vals[cmdi].(*sql.RawBytes); ok {
			str := string(*bytes)
			return &str
		} else {
			panic("not raw bytes")
		}
	}

	return nil
}

type selectQuery string

type procEntry struct {
	Id    int64
	Time  int
	Query selectQuery
}

func (qry *selectQuery) c14n() string {
	out := string(*qry)

	out = regexp.MustCompile(`"(?:\\"|""|[^"])+"|'(?:\\'|''|[^'])+'`).ReplaceAllString(out, "[[string]]")

	// @todo negative numbers present interesting problems

	lastOut := out
	for { //solves a problem with sets like 10,20,30 when there are no lookaround options as in go
		out = regexp.MustCompile(`(?m)(^|\s|,|\()\d+\.\d+($|\s|,|\))`).ReplaceAllString(out, `$1[[float]]$2`)
		if out == lastOut {
			break
		}
		lastOut = out
	}

	lastOut = out
	for {
		out = regexp.MustCompile(`(?m)(^|\s|,|\()\-?\d+($|\s|,|\))`).ReplaceAllString(out, `$1[[int]]$2`)
		if out == lastOut {
			break
		}
		lastOut = out
	}

	out = regexp.MustCompile(`\((?:\s*\[\[([a-z]+)\]\]\s*,?\s*)+\)`).ReplaceAllString(out, `[[$1-list]]`)

	return out
}

func (qry *selectQuery) csha1() string {
	h := sha1.New()
	io.WriteString(h, qry.c14n())
	return fmt.Sprintf("%x", h.Sum(nil))
}

type explainEntry struct {
	Table string
	Key   string
	Rows  int
}

func (qry *selectQuery) explain(db *sql.DB) ([]explainEntry, error) {
	output := make([]explainEntry, 0)

	rows, err := db.Query("EXPLAIN " + string(*qry))
	if err != nil {
		return output, fmt.Errorf("Explain Error, %s", err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		panic(err)
	}

	for rows.Next() {
		vals := make([]interface{}, len(cols))
		for i := range cols {
			vals[i] = new(sql.RawBytes)
		}
		err = rows.Scan(vals...)
		if err != nil {
			panic(err)
		}

		tbl := getColByName("table", cols, vals)
		rows := getColByName("rows", cols, vals)
		key := getColByName("key", cols, vals)

		rowInt, err := strconv.Atoi(*rows)
		if err != nil {
			rowInt = 0
		}

		output = append(output, explainEntry{
			Table: *tbl,
			Rows:  rowInt,
			Key:   *key,
		})
	}

	return output, nil
}

func getActiveQueries(db *sql.DB) []procEntry {
	rows, err := db.Query("SHOW FULL PROCESSLIST")
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		panic(err)
	}

	output := make([]procEntry, 0)

	isSelect := regexp.MustCompile("(?i)^\\s*select\\s")

	for rows.Next() {
		vals := make([]interface{}, len(cols))
		for i := range cols {
			vals[i] = new(sql.RawBytes)
		}
		err = rows.Scan(vals...)
		if err != nil {
			panic(err)
		}

		id := getColByName("Id", cols, vals)
		timez := getColByName("Time", cols, vals)
		cmd := getColByName("Command", cols, vals)
		info := getColByName("Info", cols, vals)

		if *cmd != "Query" || !isSelect.MatchString(*info) {
			continue
		}

		idInt, err := strconv.ParseInt(*id, 10, 64)
		if err != nil {
			idInt = int64(0)
		}

		timeInt, err := strconv.Atoi(*timez)
		if err != nil {
			panic(err)
		}

		output = append(output, procEntry{
			Id:    idInt,
			Time:  timeInt,
			Query: selectQuery(*info),
		})
	}

	return output
}
<commit_msg>No need to continuously recompile regex<commit_after>package main

import (
	"crypto/sha1"
	"database/sql"
	"fmt"
	"io"
	"log"
	"regexp"
	"strconv"

	_ "github.com/go-sql-driver/mysql"
)

func colPos(slice []string, value string) int {
	for p, v := range slice {
		if v == value {
			return p
		}
	}

	return -1
}

func getColByName(name string, cols []string, vals []interface{}) *string {
	if cmdi := colPos(cols, name); cmdi != -1 {
		if bytes, ok := vals[cmdi].(*sql.RawBytes); ok {
			str := string(*bytes)
			return &str
		} else {
			panic("not raw bytes")
		}
	}

	return nil
}

type selectQuery string

type procEntry struct {
	Id    int64
	Time  int
	Query selectQuery
}

func (qry *selectQuery) c14n() string {
	out := string(*qry)

	out = regexp.MustCompile(`"(?:\\"|""|[^"])+"|'(?:\\'|''|[^'])+'`).ReplaceAllString(out, "[[string]]")

	// @todo negative numbers present interesting problems

	lastOut := out
	for { //solves a problem with sets like 10,20,30 when there are no lookaround options as in go
		out = regexp.MustCompile(`(?m)(^|\s|,|\()\d+\.\d+($|\s|,|\))`).ReplaceAllString(out, `$1[[float]]$2`)
		if out == lastOut {
			break
		}
		lastOut = out
	}

	lastOut = out
	for {
		out = regexp.MustCompile(`(?m)(^|\s|,|\()\-?\d+($|\s|,|\))`).ReplaceAllString(out, `$1[[int]]$2`)
		if out == lastOut {
			break
		}
		lastOut = out
	}

	out = regexp.MustCompile(`\((?:\s*\[\[([a-z]+)\]\]\s*,?\s*)+\)`).ReplaceAllString(out, `[[$1-list]]`)

	return out
}

func (qry *selectQuery) csha1() string {
	h := sha1.New()
	io.WriteString(h, qry.c14n())
	return fmt.Sprintf("%x", h.Sum(nil))
}

type explainEntry struct {
	Table string
	Key   string
	Rows  int
}

func (qry *selectQuery) explain(db *sql.DB) ([]explainEntry, error) {
	output := make([]explainEntry, 0)

	rows, err := db.Query("EXPLAIN " + string(*qry))
	if err != nil {
		return output, fmt.Errorf("Explain Error, %s", err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		panic(err)
	}

	for rows.Next() {
		vals := make([]interface{}, len(cols))
		for i := range cols {
			vals[i] = new(sql.RawBytes)
		}
		err = rows.Scan(vals...)
		if err != nil {
			panic(err)
		}

		tbl := getColByName("table", cols, vals)
		rows := getColByName("rows", cols, vals)
		key := getColByName("key", cols, vals)

		rowInt, err := strconv.Atoi(*rows)
		if err != nil {
			rowInt = 0
		}

		output = append(output, explainEntry{
			Table: *tbl,
			Rows:  rowInt,
			Key:   *key,
		})
	}

	return output, nil
}

var isSelect = regexp.MustCompile("(?i)^\\s*select\\s")

func getActiveQueries(db *sql.DB) []procEntry {
	rows, err := db.Query("SHOW FULL PROCESSLIST")
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		panic(err)
	}

	output := make([]procEntry, 0)

	for rows.Next() {
		vals := make([]interface{}, len(cols))
		for i := range cols {
			vals[i] = new(sql.RawBytes)
		}
		err = rows.Scan(vals...)
		if err != nil {
			panic(err)
		}

		id := getColByName("Id", cols, vals)
		timez := getColByName("Time", cols, vals)
		cmd := getColByName("Command", cols, vals)
		info := getColByName("Info", cols, vals)

		if *cmd != "Query" || !isSelect.MatchString(*info) {
			continue
		}

		idInt, err := strconv.ParseInt(*id, 10, 64)
		if err != nil {
			idInt = int64(0)
		}

		timeInt, err := strconv.Atoi(*timez)
		if err != nil {
			panic(err)
		}

		output = append(output, procEntry{
			Id:    idInt,
			Time:  timeInt,
			Query: selectQuery(*info),
		})
	}

	return output
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"errors"
	"log"

	_ "code.google.com/p/go-sqlite/go1/sqlite3"
	"github.com/NSkelsey/btcsubprotos/ahimsa"
	"github.com/NSkelsey/protocol"
	"github.com/conformal/btcwire"
)

var (
	errNoDb = errors.New("Could not find a db to load")
)

type LiteDb struct {
	writes int
	conn   *sql.DB
}

type blockRecord struct {
	// maps to a row stored in the db
	hash     *btcwire.ShaHash
	prevhash *btcwire.ShaHash
	height   int
}

func LoadDb(dbpath string) (*LiteDb, error) {
	conn, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		log.Fatal(err)
	}

	db := &LiteDb{
		conn: conn,
	}

	return db, nil
}

func (db *LiteDb) CurrentHeight() int64 {
	// Returns the current height of the blocks in the db, if db is not initialized
	// return 0.
	cmd := `SELECT max(height) FROM blocks`
	rows, err := db.conn.Query(cmd)
	if err != nil {
		return 0
	}
	defer rows.Close()

	rows.Next()
	var height int64
	err = rows.Scan(&height)
	if err != nil {
		log.Println(err)
		return 0
	}
	return height
}

func InitDb(dbpath string) (*LiteDb, error) {
	conn, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		return nil, err
	}

	// Get the database schema for the public record.
	create, err := protocol.GetCreateSql()
	if err != nil {
		return nil, err
	}

	dropcmd := `
	DROP TABLE IF EXISTS blocks;
	DROP TABLE IF EXISTS bulletins;
	`

	// DROP db if it exists and recreate it.
	_, err = conn.Exec(dropcmd + create)
	if err != nil {
		return nil, err
	}

	db := &LiteDb{
		conn: conn,
	}

	return db, nil
}

func (db *LiteDb) storeBlockHead(bh *btcwire.BlockHeader, height int) error {
	// Writes a block to the sqlite db

	cmd := `INSERT INTO blocks (hash, prevhash, height) VALUES($1, $2, $3)`

	hash, _ := bh.BlockSha()

	println(hash.String(), height)
	_, err := db.conn.Exec(cmd, hash.String(), bh.PrevBlock.String(), height)
	if err != nil {
		return err
	}
	return nil
}

func (db *LiteDb) storeBulletin(bltn *ahimsa.Bulletin) error {
	// Writes a bulletin into the sqlite db, runs an insert or update

	var err error
	if bltn.Block == nil {
		cmd := `INSERT OR REPLACE INTO bulletins (txid, author, topic, message) VALUES($1, $2, $3, $4)`
		_, err = db.conn.Exec(cmd,
			bltn.Txid.String(),
			bltn.Author,
			bltn.Topic,
			bltn.Message,
		)
	} else {
		blockstr := bltn.Block.String()
		cmd := `INSERT OR REPLACE INTO bulletins (txid, block, author, topic, message) VALUES($1, $2, $3, $4, $5)`
		_, err = db.conn.Exec(cmd,
			bltn.Txid.String(),
			blockstr,
			bltn.Author,
			bltn.Topic,
			bltn.Message,
		)
	}
	if err != nil {
		return err
	}

	return nil
}

func (db *LiteDb) BatchInsertBH(blcks []*Block, height int) error {

	stmt, err := db.conn.Prepare("INSERT INTO blocks (hash, prevhash, height) VALUES(?, ?, ?)")
	if err != nil {
		return err
	}

	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}

	for _, blk := range blcks {
		bh := btcBHFromBH(*blk.Head)
		hash, _ := bh.BlockSha()
		prevh := bh.PrevBlock
		_, err = tx.Stmt(stmt).Exec(hash.String(), prevh.String(), height-blk.depth)
		if err != nil {
			return err
		}
	}
	err = tx.Commit()
	if err != nil {
		return err
	}

	return nil
}

func (db *LiteDb) GetBlkRecord(target *btcwire.ShaHash) (*blockRecord, error) {
	cmd := `SELECT hash, prevhash, height FROM blocks WHERE hash=$1`
	rows, err := db.conn.Query(cmd, target.String())
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	blkrec, err := scanBlkRec(rows)
	if err != nil {
		return nil, err
	}
	return blkrec, nil
}

func (db *LiteDb) GetChainTip() (*blockRecord, error) {
	cmd := `SELECT hash, prevhash, max(height) FROM blocks`
	rows, err := db.conn.Query(cmd)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	blkrec, err := scanBlkRec(rows)
	if err != nil {
		return nil, err
	}
	return blkrec, nil
}

func scanBlkRec(rows *sql.Rows) (*blockRecord, error) {
	rows.Next()
	// called for effect
	var hash, prevhash string
	var height int
	if err := rows.Scan(&hash, &prevhash, &height); err != nil {
		return nil, err
	}

	btchash, err := btcwire.NewShaHashFromStr(hash)
	if err != nil {
		return nil, err
	}

	btcprevhash, err := btcwire.NewShaHashFromStr(prevhash)
	if err != nil {
		return nil, err
	}

	blkrec := &blockRecord{
		hash:     btchash,
		prevhash: btcprevhash,
		height:   height,
	}
	return blkrec, nil
}
<commit_msg>import reorg<commit_after>package main

import (
	"database/sql"
	"errors"
	"log"

	_ "code.google.com/p/go-sqlite/go1/sqlite3"
	"github.com/NSkelsey/protocol/ahimsa"
	"github.com/conformal/btcwire"
)

var (
	errNoDb = errors.New("Could not find a db to load")
)

type LiteDb struct {
	writes int
	conn   *sql.DB
}

type blockRecord struct {
	// maps to a row stored in the db
	hash     *btcwire.ShaHash
	prevhash *btcwire.ShaHash
	height   int
}

func LoadDb(dbpath string) (*LiteDb, error) {
	conn, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		log.Fatal(err)
	}

	db := &LiteDb{
		conn: conn,
	}

	return db, nil
}

func (db *LiteDb) CurrentHeight() int64 {
	// Returns the current height of the blocks in the db, if db is not initialized
	// return 0.
	cmd := `SELECT max(height) FROM blocks`
	rows, err := db.conn.Query(cmd)
	if err != nil {
		return 0
	}
	defer rows.Close()

	rows.Next()
	var height int64
	err = rows.Scan(&height)
	if err != nil {
		log.Println(err)
		return 0
	}
	return height
}

func InitDb(dbpath string) (*LiteDb, error) {
	conn, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		return nil, err
	}

	// Get the database schema for the public record.
	create, err := ahimsa.GetCreateSql()
	if err != nil {
		return nil, err
	}

	dropcmd := `
	DROP TABLE IF EXISTS blocks;
	DROP TABLE IF EXISTS bulletins;
	`

	// DROP db if it exists and recreate it.
	_, err = conn.Exec(dropcmd + create)
	if err != nil {
		return nil, err
	}

	db := &LiteDb{
		conn: conn,
	}

	return db, nil
}

func (db *LiteDb) storeBlockHead(bh *btcwire.BlockHeader, height int) error {
	// Writes a block to the sqlite db

	cmd := `INSERT INTO blocks (hash, prevhash, height) VALUES($1, $2, $3)`

	hash, _ := bh.BlockSha()

	println(hash.String(), height)
	_, err := db.conn.Exec(cmd, hash.String(), bh.PrevBlock.String(), height)
	if err != nil {
		return err
	}
	return nil
}

func (db *LiteDb) storeBulletin(bltn *ahimsa.Bulletin) error {
	// Writes a bulletin into the sqlite db, runs an insert or update

	var err error
	if bltn.Block == nil {
		cmd := `INSERT OR REPLACE INTO bulletins (txid, author, topic, message) VALUES($1, $2, $3, $4)`
		_, err = db.conn.Exec(cmd,
			bltn.Txid.String(),
			bltn.Author,
			bltn.Topic,
			bltn.Message,
		)
	} else {
		blockstr := bltn.Block.String()
		cmd := `INSERT OR REPLACE INTO bulletins (txid, block, author, topic, message) VALUES($1, $2, $3, $4, $5)`
		_, err = db.conn.Exec(cmd,
			bltn.Txid.String(),
			blockstr,
			bltn.Author,
			bltn.Topic,
			bltn.Message,
		)
	}
	if err != nil {
		return err
	}

	return nil
}

func (db *LiteDb) BatchInsertBH(blcks []*Block, height int) error {

	stmt, err := db.conn.Prepare("INSERT INTO blocks (hash, prevhash, height) VALUES(?, ?, ?)")
	if err != nil {
		return err
	}

	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}

	for _, blk := range blcks {
		bh := btcBHFromBH(*blk.Head)
		hash, _ := bh.BlockSha()
		prevh := bh.PrevBlock
		_, err = tx.Stmt(stmt).Exec(hash.String(), prevh.String(), height-blk.depth)
		if err != nil {
			return err
		}
	}
	err = tx.Commit()
	if err != nil {
		return err
	}

	return nil
}

func (db *LiteDb) GetBlkRecord(target *btcwire.ShaHash) (*blockRecord, error) {
	cmd := `SELECT hash, prevhash, height FROM blocks WHERE hash=$1`
	rows, err := db.conn.Query(cmd, target.String())
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	blkrec, err := scanBlkRec(rows)
	if err != nil {
		return nil, err
	}
	return blkrec, nil
}

func (db *LiteDb) GetChainTip() (*blockRecord, error) {
	cmd := `SELECT hash, prevhash, max(height) FROM blocks`
	rows, err := db.conn.Query(cmd)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	blkrec, err := scanBlkRec(rows)
	if err != nil {
		return nil, err
	}
	return blkrec, nil
}

func scanBlkRec(rows *sql.Rows) (*blockRecord, error) {
	rows.Next()
	// called for effect
	var hash, prevhash string
	var height int
	if err := rows.Scan(&hash, &prevhash, &height); err != nil {
		return nil, err
	}

	btchash, err := btcwire.NewShaHashFromStr(hash)
	if err != nil {
		return nil, err
	}

	btcprevhash, err := btcwire.NewShaHashFromStr(prevhash)
	if err != nil {
		return nil, err
	}

	blkrec := &blockRecord{
		hash:     btchash,
		prevhash: btcprevhash,
		height:   height,
	}
	return blkrec, nil
}
<|endoftext|>


<commit_before>package cli_test

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/alecthomas/template"
	"github.com/wendigo/go-bind-plugin/cli"
)

type testCase struct {
	Plugin         string
	ExpectedOutput string
	ExecutedCode   string
	AsInterface    bool
}

func TestWillGenerateComplexPluginWithoutErrors(t *testing.T) {

	testCases := []testCase{
		{
			Plugin:         "complex_plugin",
			ExecutedCode:   "fmt.Println(pl)",
			ExpectedOutput: "",
		},
		{
			Plugin:         "basic_plugin",
			ExecutedCode:   "fmt.Println(pl.ReturningIntArray())",
			ExpectedOutput: "[1 0 1]",
		},
		{
			Plugin:         "plugin_as_interface",
			ExecutedCode:   "fmt.Println(pl.ReturningStringSlice())",
			ExpectedOutput: "hello world",
			AsInterface:    true,
		},
	}

	for i, testCase := range testCases {
		t.Logf("[Test %d] Generating %s plugin...", i, testCase.Plugin)

		config := cli.Config{
			PluginPackage:      fmt.Sprintf("./internal/test_fixtures/%s", testCase.Plugin),
			OutputPath:         fmt.Sprintf("./internal/test_fixtures/generated/%s/plugin.go", testCase.Plugin),
			PluginPath:         fmt.Sprintf("./internal/test_fixtures/generated/%s/plugin.so", testCase.Plugin),
			FormatCode:         true,
			CheckSha256:        true,
			ForcePluginRebuild: true,
			OutputPackage:      "main",
			OutputName:         "TestWrapper",
			AsInterface:        testCase.AsInterface,
		}

		t.Logf("[Test %d] Generator config: %+v", i, config)

		if err := generatePluginWithCli(config, t); err != nil {
			t.Fatalf("[Test %d] Expected error to be nil, actual: %s", i, err)
		}

		runFile := fmt.Sprintf("./internal/test_fixtures/generated/%s/plugin.go", testCase.Plugin)

		t.Logf("[Test %d] Running plugin via %s", i, runFile)
		output, err := runPlugin(testCase.ExecutedCode, runFile, config)

		if err != nil {
			t.Fatalf("[Test %d] Expected err to be nil, actual: %s", i, err)
		}

		if !strings.Contains(output, testCase.ExpectedOutput) {
			t.Fatalf("[Test %d] Expected output to contain %s, actual output:\n=======\n%s\n=======\n", i, testCase.ExpectedOutput, output)
		}
	}
}

// Switch to generatePluginWithCli when https://github.com/golang/go/issues/17928 is solved
func generatePluginWithCli(config cli.Config, t *testing.T) error {
	client, err := cli.New(config, log.New(os.Stdout, "", 0))

	if err != nil {
		return err
	}

	if generateErr := client.GenerateFile(); generateErr != nil {
		return generateErr
	}

	return nil
}

func generatePluginViaCommandLine(config cli.Config, t *testing.T) error {

	args := []string{"run", "../main.go"}
	args = append(args, strings.Split(config.String(), " ")...)
	cmd := exec.Command("go", args...)

	t.Logf("Generating plugin with config: %+v", cmd)

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	return cmd.Run()
}

func runPlugin(code string, path string, config cli.Config) (string, error) {
	file, err := os.OpenFile(config.OutputPath, os.O_APPEND|os.O_WRONLY, 0700)

	if err != nil {
		return "", err
	}

	tmp, err := template.New("test_case").Parse(runTemplate)
	if err != nil {
		return "", err
	}

	if err := tmp.Execute(file, struct {
		Config cli.Config
		Code   string
	}{
		Config: config,
		Code:   code,
	}); err != nil {
		return "", err
	}

	var outBuffer bytes.Buffer

	cmd := exec.Command("go", "run", config.OutputPath)
	cmd.Stdout = bufio.NewWriter(&outBuffer)
	cmd.Stderr = os.Stderr

	runErr := cmd.Run()
	return string(outBuffer.Bytes()), runErr
}

var runTemplate = `

func main() {
  pl, err := Bind{{.Config.OutputName}}("{{.Config.PluginPath}}")

  if err != nil {
    fmt.Println(err)
    os.Exit(1)
  }

  {{.Code}}
}
`
<commit_msg>Dereference variables in test<commit_after>package cli_test

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/alecthomas/template"
	"github.com/wendigo/go-bind-plugin/cli"
)

type testCase struct {
	Plugin         string
	ExpectedOutput string
	ExecutedCode   string
	AsInterface    bool
}

func TestWillGenerateComplexPluginWithoutErrors(t *testing.T) {

	testCases := []testCase{
		{
			Plugin:         "complex_plugin",
			ExecutedCode:   "fmt.Println(pl)",
			ExpectedOutput: "",
		},
		{
			Plugin:         "basic_plugin",
			ExecutedCode:   "fmt.Println(pl.ReturningIntArray())",
			ExpectedOutput: "[1 0 1]",
		},
		{
			Plugin:         "plugin_as_interface",
			ExecutedCode:   "fmt.Println(pl.ReturningStringSlice())",
			ExpectedOutput: "hello world",
			AsInterface:    true,
		},
	}

	for i, testCase := range testCases {
		t.Logf("[Test %d] Generating %s plugin...", i, testCase.Plugin)

		config := cli.Config{
			PluginPackage:        fmt.Sprintf("./internal/test_fixtures/%s", testCase.Plugin),
			OutputPath:           fmt.Sprintf("./internal/test_fixtures/generated/%s/plugin.go", testCase.Plugin),
			PluginPath:           fmt.Sprintf("./internal/test_fixtures/generated/%s/plugin.so", testCase.Plugin),
			FormatCode:           true,
			CheckSha256:          true,
			ForcePluginRebuild:   true,
			OutputPackage:        "main",
			OutputName:           "TestWrapper",
			AsInterface:          testCase.AsInterface,
			DereferenceVariables: true,
		}

		t.Logf("[Test %d] Generator config: %+v", i, config)

		if err := generatePluginWithCli(config, t); err != nil {
			t.Fatalf("[Test %d] Expected error to be nil, actual: %s", i, err)
		}

		runFile := fmt.Sprintf("./internal/test_fixtures/generated/%s/plugin.go", testCase.Plugin)

		t.Logf("[Test %d] Running plugin via %s", i, runFile)
		output, err := runPlugin(testCase.ExecutedCode, runFile, config)

		if err != nil {
			t.Fatalf("[Test %d] Expected err to be nil, actual: %s", i, err)
		}

		if !strings.Contains(output, testCase.ExpectedOutput) {
			t.Fatalf("[Test %d] Expected output to contain %s, actual output:\n=======\n%s\n=======\n", i, testCase.ExpectedOutput, output)
		}
	}
}

// Switch to generatePluginWithCli when https://github.com/golang/go/issues/17928 is solved
func generatePluginWithCli(config cli.Config, t *testing.T) error {
	client, err := cli.New(config, log.New(os.Stdout, "", 0))

	if err != nil {
		return err
	}

	if generateErr := client.GenerateFile(); generateErr != nil {
		return generateErr
	}

	return nil
}

func generatePluginViaCommandLine(config cli.Config, t *testing.T) error {

	args := []string{"run", "../main.go"}
	args = append(args, strings.Split(config.String(), " ")...)
	cmd := exec.Command("go", args...)

	t.Logf("Generating plugin with config: %+v", cmd)

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	return cmd.Run()
}

func runPlugin(code string, path string, config cli.Config) (string, error) {
	file, err := os.OpenFile(config.OutputPath, os.O_APPEND|os.O_WRONLY, 0700)

	if err != nil {
		return "", err
	}

	tmp, err := template.New("test_case").Parse(runTemplate)
	if err != nil {
		return "", err
	}

	if err := tmp.Execute(file, struct {
		Config cli.Config
		Code   string
	}{
		Config: config,
		Code:   code,
	}); err != nil {
		return "", err
	}

	var outBuffer bytes.Buffer

	cmd := exec.Command("go", "run", config.OutputPath)
	cmd.Stdout = bufio.NewWriter(&outBuffer)
	cmd.Stderr = os.Stderr

	runErr := cmd.Run()
	return string(outBuffer.Bytes()), runErr
}

var runTemplate = `

func main() {
  pl, err := Bind{{.Config.OutputName}}("{{.Config.PluginPath}}")

  if err != nil {
    fmt.Println(err)
    os.Exit(1)
  }

  {{.Code}}
}
`
<|endoftext|>


<commit_before>package cli

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

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-utils/colorstring"
	"github.com/codegangsta/cli"
)

// ValidationItemModel ...
type ValidationItemModel struct {
	IsValid bool   `json:"is_valid" yaml:"is_valid"`
	Error   string `json:"error,omitempty" yaml:"error,omitempty"`
}

// ValidationModel ...
type ValidationModel struct {
	Config  *ValidationItemModel `json:"config,omitempty" yaml:"config,omitempty"`
	Secrets *ValidationItemModel `json:"secrets,omitempty" yaml:"secrets,omitempty"`
}

func printRawValidation(validation ValidationModel) error {
	validConfig := true
	if validation.Config != nil {
		fmt.Println(colorstring.Blue("Config validation result:"))
		configValidation := *validation.Config
		if configValidation.IsValid {
			fmt.Printf("is valid: %s\n", colorstring.Greenf("%v", configValidation.IsValid))
		} else {
			fmt.Printf("is valid: %s\n", colorstring.Redf("%v", configValidation.IsValid))
			fmt.Printf("error: %s\n", colorstring.Red(configValidation.Error))

			validConfig = false
		}
		fmt.Println()
	}

	validSecrets := true
	if validation.Secrets != nil {
		fmt.Println(colorstring.Blue("Secret validation result:"))
		secretValidation := *validation.Secrets
		if secretValidation.IsValid {
			fmt.Printf("is valid: %s\n", colorstring.Greenf("%v", secretValidation.IsValid))
		} else {
			fmt.Printf("is valid: %s\n", colorstring.Redf("%v", secretValidation.IsValid))
			fmt.Printf("error: %s\n", colorstring.Red(secretValidation.Error))

			validSecrets = false
		}
	}

	if !validConfig && !validSecrets {
		return errors.New("Config and secrets are invalid")
	} else if !validConfig {
		return errors.New("Config is invalid")
	} else if !validSecrets {
		return errors.New("Secret is invalid")
	}
	return nil
}

func printJSONValidation(validation ValidationModel) error {
	bytes, err := json.Marshal(validation)
	if err != nil {
		return err
	}

	fmt.Println(string(bytes))
	if (validation.Config != nil && !validation.Config.IsValid) &&
		(validation.Secrets != nil && !validation.Secrets.IsValid) {
		return errors.New("Config and secrets are invalid")
	} else if validation.Config != nil && !validation.Config.IsValid {
		return errors.New("Config is invalid")
	} else if validation.Secrets != nil && !validation.Secrets.IsValid {
		return errors.New("Secret is invalid")
	}
	return nil
}

func validate(c *cli.Context) {
	format := c.String(OuputFormatKey)
	if format == "" {
		format = OutputFormatRaw
	} else if !(format == OutputFormatRaw || format == OutputFormatJSON) {
		registerFatal(fmt.Sprintf("Invalid format: %s", format), OutputFormatJSON)
	}

	validation := ValidationModel{}

	pth, err := GetBitriseConfigFilePath(c)
	if err != nil && err.Error() != "No workflow yml found" {
		registerFatal(fmt.Sprintf("Faild to get config path, err: %s", err), format)
	}
	if pth != "" || (pth == "" && c.String(ConfigBase64Key) != "") {
		// Config validation
		isValid := true
		errMsg := ""

		_, err := CreateBitriseConfigFromCLIParams(c)
		if err != nil {
			isValid = false
			errMsg = err.Error()
		}

		validation.Config = &ValidationItemModel{
			IsValid: isValid,
			Error:   errMsg,
		}
	} else {
		log.Debug("No config found for validation")
	}

	pth, err = GetInventoryFilePath(c)
	if err != nil {
		registerFatal(fmt.Sprintf("Faild to get secrets path, err: %s", err), format)
	}
	if pth != "" || c.String(InventoryBase64Key) != "" {
		// Inventory validation
		isValid := true
		errMsg := ""

		_, err := CreateInventoryFromCLIParams(c)
		if err != nil {
			isValid = false
			errMsg = err.Error()
		}

		validation.Secrets = &ValidationItemModel{
			IsValid: isValid,
			Error:   errMsg,
		}
	}

	if validation.Config == nil && validation.Secrets == nil {
		registerFatal("No config or secrets found for validation", format)
	}

	switch format {
	case OutputFormatRaw:
		if err := printRawValidation(validation); err != nil {
			registerFatal(fmt.Sprintf("Validation failed, err: %s", err), format)
		}
		break
	case OutputFormatJSON:
		if err := printJSONValidation(validation); err != nil {
			registerFatal(fmt.Sprintf("Validation failed, err: %s", err), format)
		}
		break
	default:
		registerFatal(fmt.Sprintf("Invalid format: %s", format), OutputFormatJSON)
	}
}
<commit_msg>validate fix<commit_after>package cli

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

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-utils/colorstring"
	"github.com/codegangsta/cli"
)

// ValidationItemModel ...
type ValidationItemModel struct {
	IsValid bool   `json:"is_valid" yaml:"is_valid"`
	Error   string `json:"error,omitempty" yaml:"error,omitempty"`
}

// ValidationModel ...
type ValidationModel struct {
	Config  *ValidationItemModel `json:"config,omitempty" yaml:"config,omitempty"`
	Secrets *ValidationItemModel `json:"secrets,omitempty" yaml:"secrets,omitempty"`
}

func printRawValidation(validation ValidationModel) error {
	validConfig := true
	if validation.Config != nil {
		fmt.Println(colorstring.Blue("Config validation result:"))
		configValidation := *validation.Config
		if configValidation.IsValid {
			fmt.Printf("is valid: %s\n", colorstring.Greenf("%v", configValidation.IsValid))
		} else {
			fmt.Printf("is valid: %s\n", colorstring.Redf("%v", configValidation.IsValid))
			fmt.Printf("error: %s\n", colorstring.Red(configValidation.Error))

			validConfig = false
		}
		fmt.Println()
	}

	validSecrets := true
	if validation.Secrets != nil {
		fmt.Println(colorstring.Blue("Secret validation result:"))
		secretValidation := *validation.Secrets
		if secretValidation.IsValid {
			fmt.Printf("is valid: %s\n", colorstring.Greenf("%v", secretValidation.IsValid))
		} else {
			fmt.Printf("is valid: %s\n", colorstring.Redf("%v", secretValidation.IsValid))
			fmt.Printf("error: %s\n", colorstring.Red(secretValidation.Error))

			validSecrets = false
		}
	}

	if !validConfig && !validSecrets {
		return errors.New("Config and secrets are invalid")
	} else if !validConfig {
		return errors.New("Config is invalid")
	} else if !validSecrets {
		return errors.New("Secret is invalid")
	}
	return nil
}

func printJSONValidation(validation ValidationModel) {
	bytes, err := json.Marshal(validation)
	if err != nil {
		registerFatal(fmt.Sprintf("Failed to parse validation result, err: %s, result: %#v", err, validation), OutputFormatJSON)
	}

	fmt.Println(string(bytes))
}

func validate(c *cli.Context) {
	format := c.String(OuputFormatKey)
	if format == "" {
		format = OutputFormatRaw
	} else if !(format == OutputFormatRaw || format == OutputFormatJSON) {
		registerFatal(fmt.Sprintf("Invalid format: %s", format), OutputFormatJSON)
	}

	validation := ValidationModel{}

	pth, err := GetBitriseConfigFilePath(c)
	if err != nil && err.Error() != "No workflow yml found" {
		registerFatal(fmt.Sprintf("Faild to get config path, err: %s", err), format)
	}
	if pth != "" || (pth == "" && c.String(ConfigBase64Key) != "") {
		// Config validation
		isValid := true
		errMsg := ""

		_, err := CreateBitriseConfigFromCLIParams(c)
		if err != nil {
			isValid = false
			errMsg = err.Error()
		}

		validation.Config = &ValidationItemModel{
			IsValid: isValid,
			Error:   errMsg,
		}
	} else {
		log.Debug("No config found for validation")
	}

	pth, err = GetInventoryFilePath(c)
	if err != nil {
		registerFatal(fmt.Sprintf("Faild to get secrets path, err: %s", err), format)
	}
	if pth != "" || c.String(InventoryBase64Key) != "" {
		// Inventory validation
		isValid := true
		errMsg := ""

		_, err := CreateInventoryFromCLIParams(c)
		if err != nil {
			isValid = false
			errMsg = err.Error()
		}

		validation.Secrets = &ValidationItemModel{
			IsValid: isValid,
			Error:   errMsg,
		}
	}

	if validation.Config == nil && validation.Secrets == nil {
		registerFatal("No config or secrets found for validation", format)
	}

	switch format {
	case OutputFormatRaw:
		if err := printRawValidation(validation); err != nil {
			registerFatal(fmt.Sprintf("Validation failed, err: %s", err), format)
		}
		break
	case OutputFormatJSON:
		printJSONValidation(validation)
		break
	default:
		registerFatal(fmt.Sprintf("Invalid format: %s", format), OutputFormatJSON)
	}
}
<|endoftext|>


<commit_before>package unicornhat

// #include "ws2812-RPi.h"
// #include "unicornhat-bridge.h"
import "C"
import "fmt"

type Color struct {
	r, g, b byte
}

// initialization of hardware
func InitHardware() {
	C.initHardware()
}

func Initialize(numPixels int) {
	C.init(C.int(numPixels))
}

func StartTransfer() {
	C.startTransfer()
}

// Led updates
func Show() {
	C.show()
}

func GetBrightness() float64 {
	return float64(C.getBrightness())
}
func SetBrightness(brightness float64) byte {
	return byte(C.setBrightness(C.double(brightness)))
}

func ClearPWMBuffer() {
	C.clearPWMBuffer()
}

func Clear() {
	C.clear()
}

func ClearLEDBuffer() {
	C.clearLEDBuffer()
}

func GetPixelColor(pixel uint) Color {
	container := C.getPixelColor(C.uint(pixel))
	return NewColor(byte(container.r), byte(container.g), byte(container.b))
}

func NewColor(r, g, b byte) Color {
	return Color { r: r, g: g, b: b }
}
<commit_msg>Basic interface library implemented<commit_after>// Interface functions with the pimoroni UnicornHat
package unicornhat

//#include "ws2812-RPi.h"
//
//double getDefaultBrightnessMacro() {
//	return DEFAULT_BRIGHTNESS;
//}
import "C"

type Pixel struct {
	R, G, B byte
}

// initialization of hardware
func InitHardware() {
	C.initHardware()
}

func Initialize(numPixels int) {
	C.init(C.int(numPixels))
}

func StartTransfer() {
	C.startTransfer()
}

// Led updates
func Show() {
	C.show()
}

func GetBrightness() float64 {
	return float64(C.getBrightness())
}
func SetBrightness(brightness float64) bool {
	return (C.setBrightness(C.double(brightness)) == 1)
}

func ClearPWMBuffer() {
	C.clearPWMBuffer()
}

func Clear() {
	C.clear()
}

func ClearLEDBuffer() {
	C.clearLEDBuffer()
}

func GetPixelColor(pixel uint) Pixel {
	return fromNativePixel(C.getPixelColor(C.uint(pixel)))
}
func NewPixel(r, g, b byte) Pixel {
	return Pixel { R: r, G: g, B: b }

}
func Color(r, g, b byte) Pixel {
	return fromNativePixel(C.Color(C.uchar(r), C.uchar(g), C.uchar(b)))
}
func RGB2Color(r, g, b byte) Pixel {
	return fromNativePixel(C.RGB2Color(C.uchar(r), C.uchar(g), C.uchar(b)))
}

func (this Pixel) nativePixel() C.Color_t {
	var v C.Color_t
	v.r = C.uchar(this.R)
	v.g = C.uchar(this.G)
	v.b = C.uchar(this.B)
	return v
}

func fromNativePixel(pixel C.Color_t) Pixel {
	return NewPixel(byte(pixel.r), byte(pixel.g), byte(pixel.b))
}

func NumPixels() uint {
	return uint(C.numPixels())
}

func SetPixelColor(pixel uint, r, g, b byte) bool {
	return (C.setPixelColor(C.uint(pixel), C.uchar(r), C.uchar(g), C.uchar(b)) == 1)
}

func SetPixelColorType(pixel uint, color Pixel) bool {
	return (C.setPixelColorT(C.uint(pixel), color.nativePixel()) == 1)
}

func DefaultBrightness() float64 {
	return float64(C.getDefaultBrightnessMacro())
}

<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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
//
//     https://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 main

import (
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

type localFunctionServer struct {
	output string
	cmd    string
}

func (l *localFunctionServer) Start() (func(), error) {
	args := strings.Fields(l.cmd)
	cmd := exec.Command(args[0], args[1:]...)

	stdout, err := os.Create(stdoutFile)
	if err != nil {
		return nil, err
	}
	cmd.Stdout = stdout

	stderr, err := os.Create(stderrFile)
	if err != nil {
		return nil, err
	}
	cmd.Stderr = stderr
	err = cmd.Start()
	if err != nil {
		return nil, err
	}
	log.Printf("Framework server started.")

	// Give it some time to do its setup.
	time.Sleep(time.Duration(*startDelay) * time.Second)

	shutdown := func() {
		// TODO: kill processes properly.
		if err := cmd.Process.Kill(); err != nil {
			log.Fatalf("failed to kill process: %v", err)
		}
		stdout.Close()
		stderr.Close()
		log.Printf("Framework server shut down. Wrote logs to %v and %v.", stdoutFile, stderrFile)
	}
	return shutdown, nil
}

func (l *localFunctionServer) OutputFile() ([]byte, error) {
	return ioutil.ReadFile(l.output)
}
<commit_msg>fix: properly kill child processes (#67)<commit_after>// Copyright 2020 Google LLC
//
// 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
//
//     https://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 main

import (
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

type localFunctionServer struct {
	output string
	cmd    string
}

func (l *localFunctionServer) Start() (func(), error) {
	args := strings.Fields(l.cmd)
	cmd := exec.Command(args[0], args[1:]...)

	// Set a process group ID so that later we can kill child processes too. As an
	// example, if the command is `go run main.go`, Go will build a binary in a
	// temp dir and then execute it. If we simply cmd.Process.Kill() the exec.Command
	// then the running binary will not be killed. Only if we make a group and then
	// kill the group will child processes be killed.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	stdout, err := os.Create(stdoutFile)
	if err != nil {
		return nil, err
	}
	cmd.Stdout = stdout

	stderr, err := os.Create(stderrFile)
	if err != nil {
		return nil, err
	}
	cmd.Stderr = stderr
	err = cmd.Start()
	if err != nil {
		return nil, err
	}
	log.Printf("Framework server started.")

	// Give it some time to do its setup.
	time.Sleep(time.Duration(*startDelay) * time.Second)

	shutdown := func() {
		stdout.Close()
		stderr.Close()

		pgid, err := syscall.Getpgid(cmd.Process.Pid)
		if err != nil {
			log.Printf("Failed to get pgid: %v", err)

			// Kill just the parent process since we failed to get the process group ID.
			if err := cmd.Process.Kill(); err != nil {
				log.Fatalf("Failed to kill process: %v", err)
			}
		} else {
			// Kill the whole process group.
			if err := syscall.Kill(-pgid, syscall.SIGKILL); err != nil {
				log.Fatalf("Failed to kill process group: %v", err)
			}
		}

		log.Printf("Framework server shut down. Wrote logs to %v and %v.", stdoutFile, stderrFile)
	}
	return shutdown, nil
}

func (l *localFunctionServer) OutputFile() ([]byte, error) {
	return ioutil.ReadFile(l.output)
}
<|endoftext|>


<commit_before><commit_msg>Undo debugging commit<commit_after><|endoftext|>


<commit_before><commit_msg>Remove deprecated DB interface.<commit_after><|endoftext|>