<commit_before>package gohm

import (
	"compress/flate"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/golang/snappy"
)

type compressionResponseWriter struct {
	http.ResponseWriter
	compressionWriter io.Writer
}

func (g compressionResponseWriter) Write(b []byte) (int, error) {
	return g.compressionWriter.Write(b)
}

// WithGzip returns a new http.Handler that optionally compresses the response
// text using the gzip compression algorithm when the HTTP request's
// `Accept-Encoding` header includes the string `gzip`.
//
//	mux := http.NewServeMux()
//	mux.Handle("/example/path", gohm.WithGzip(someHandler))
func WithGzip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			next.ServeHTTP(w, r)
			return
		}
		gz := gzip.NewWriter(w)
		defer func() {
			if err := gz.Close(); err != nil {
				Error(w, fmt.Sprintf("cannot compress stream: %s", err), http.StatusInternalServerError)
			}
		}()
		w.Header().Set("Content-Encoding", "gzip")
		next.ServeHTTP(compressionResponseWriter{ResponseWriter: w, compressionWriter: gz}, r)
	})
}

// WithCompression returns a new http.Handler that optionally compresses the
// response text using either the gzip or deflate compression algorithm when the
// HTTP request's `Accept-Encoding` header includes the string `gzip` or
// `deflate`. To prevent the specified next http.Handler from also seeing the
// `Accept-Encoding` request header, and possibly also compressing the data a
// second time, this function removes that header from the request.
//
//	mux := http.NewServeMux()
//	mux.Handle("/example/path", gohm.WithCompression(someHandler))
func WithCompression(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		var newWriteCloser io.WriteCloser
		var encodingName, responseHeaderName string

		requestHeaderName := "TE"
		acceptableEncodings := r.Header.Get(requestHeaderName)
		if acceptableEncodings != "" {
			// If transfer-encoding specified, then completely ignore
			// Accept-Encoding, because the upstream node has specifically
			// requested a node-to-node transfer compression algorithm.
			responseHeaderName = "Transfer-Encoding"
		} else {
			responseHeaderName = "Content-Encoding"
			requestHeaderName = "Accept-Encoding"
			acceptableEncodings = r.Header.Get(requestHeaderName)
		}

		// Because many browsers include a buggy deflate compression algorithm,
		// prefer `gzip` over `deflate` if both are acceptable.
		if encodingName = "snappy"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser = snappy.NewBufferedWriter(w)
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using snappy: %s", err), http.StatusInternalServerError)
				}
			}()
		} else if encodingName = "gzip"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser = gzip.NewWriter(w)
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using gzip: %s", err), http.StatusInternalServerError)
				}
			}()
		} else if encodingName = "deflate"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser, err = flate.NewWriter(w, -1)
			if err != nil {
				next.ServeHTTP(w, r)
				return
			}
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using deflate: %s", err), http.StatusInternalServerError)
				}
			}()
		} else {
			next.ServeHTTP(w, r)
			return
		}
		r.Header.Del(requestHeaderName)
		w.Header().Set(responseHeaderName, encodingName)
		if responseHeaderName == "Content-Encoding" {
			w.Header().Set("Vary", responseHeaderName)
		}
		next.ServeHTTP(compressionResponseWriter{ResponseWriter: w, compressionWriter: newWriteCloser}, r)
	})
}
<commit_msg>header name is title cased in comment line<commit_after>package gohm

import (
	"compress/flate"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/golang/snappy"
)

type compressionResponseWriter struct {
	http.ResponseWriter
	compressionWriter io.Writer
}

func (g compressionResponseWriter) Write(b []byte) (int, error) {
	return g.compressionWriter.Write(b)
}

// WithGzip returns a new http.Handler that optionally compresses the response
// text using the gzip compression algorithm when the HTTP request's
// `Accept-Encoding` header includes the string `gzip`.
//
//	mux := http.NewServeMux()
//	mux.Handle("/example/path", gohm.WithGzip(someHandler))
func WithGzip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			next.ServeHTTP(w, r)
			return
		}
		gz := gzip.NewWriter(w)
		defer func() {
			if err := gz.Close(); err != nil {
				Error(w, fmt.Sprintf("cannot compress stream: %s", err), http.StatusInternalServerError)
			}
		}()
		w.Header().Set("Content-Encoding", "gzip")
		next.ServeHTTP(compressionResponseWriter{ResponseWriter: w, compressionWriter: gz}, r)
	})
}

// WithCompression returns a new http.Handler that optionally compresses the
// response text using either the gzip or deflate compression algorithm when the
// HTTP request's `Accept-Encoding` header includes the string `gzip` or
// `deflate`. To prevent the specified next http.Handler from also seeing the
// `Accept-Encoding` request header, and possibly also compressing the data a
// second time, this function removes that header from the request.
//
//	mux := http.NewServeMux()
//	mux.Handle("/example/path", gohm.WithCompression(someHandler))
func WithCompression(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		var newWriteCloser io.WriteCloser
		var encodingName, responseHeaderName string

		requestHeaderName := "TE"
		acceptableEncodings := r.Header.Get(requestHeaderName)
		if acceptableEncodings != "" {
			// If Transfer-Encoding specified, then completely ignore
			// Accept-Encoding, because the upstream node has specifically
			// requested a node-to-node transfer compression algorithm.
			responseHeaderName = "Transfer-Encoding"
		} else {
			responseHeaderName = "Content-Encoding"
			requestHeaderName = "Accept-Encoding"
			acceptableEncodings = r.Header.Get(requestHeaderName)
		}

		// Because many browsers include a buggy deflate compression algorithm,
		// prefer `gzip` over `deflate` if both are acceptable.
		if encodingName = "snappy"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser = snappy.NewBufferedWriter(w)
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using snappy: %s", err), http.StatusInternalServerError)
				}
			}()
		} else if encodingName = "gzip"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser = gzip.NewWriter(w)
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using gzip: %s", err), http.StatusInternalServerError)
				}
			}()
		} else if encodingName = "deflate"; strings.Contains(acceptableEncodings, encodingName) {
			newWriteCloser, err = flate.NewWriter(w, -1)
			if err != nil {
				next.ServeHTTP(w, r)
				return
			}
			defer func() {
				if err := newWriteCloser.Close(); err != nil {
					Error(w, fmt.Sprintf("cannot compress stream using deflate: %s", err), http.StatusInternalServerError)
				}
			}()
		} else {
			next.ServeHTTP(w, r)
			return
		}
		r.Header.Del(requestHeaderName)
		w.Header().Set(responseHeaderName, encodingName)
		if responseHeaderName == "Content-Encoding" {
			w.Header().Set("Vary", responseHeaderName)
		}
		next.ServeHTTP(compressionResponseWriter{ResponseWriter: w, compressionWriter: newWriteCloser}, r)
	})
}
<|endoftext|>


<commit_before>package conf

import (
	"os"
	"os/user"
	"path/filepath"

	db "github.com/rvflash/awql-db"
	awql "github.com/rvflash/awql-driver"
	"github.com/rvflash/awql/driver"
)

// Options represents all available parameters.
type Options interface {
	AccountID() string
	APIVersion() string
	ExecuteStmt() string
	IsInteractive() bool
	SupportsZeroImpressions() bool
	UseBatchMode() bool
	UseVerboseMode() bool
	WithAutoRehash() bool
	WithCache() bool
}

// Settings represents all configurable options.
type Settings interface {
	Options
	Dsn() string
	CacheDir() string
	DatabaseDir() string
	HistoryFile() string
	Init() error
}

// Context represents the program properties.
type Context struct {
	wrkDir, homeDir string
	tk              *Credentials
	opts            *Flag
}

// New returns an instance of Credentials.
// The string dir contains the path of the current workspace.
func New(dir string) *Context {
	return &Context{wrkDir: dir, opts: Parse()}
}

// AccountID returns the Account ID.
func (c *Context) AccountID() string {
	return *c.opts.AccountID
}

// APIVersion returns the API version.
func (c *Context) APIVersion() string {
	return *c.opts.APIVersion
}

// CacheDir returns the path to store cache files.
func (c *Context) CacheDir() string {
	if c.homeDir == "" {
		return ""
	}
	return filepath.Join(c.homeDir, "cache")
}

// DatabaseDir returns the path to the database.
func (c *Context) DatabaseDir() string {
	return filepath.Join(c.wrkDir, "vendor/github.com/rvflash/awql-db/src")
}

// Dsn outputs the data source name.
func (c *Context) Dsn() string {
	// Data source name used to connect to Adwords.
	dsn := awql.NewDsn(c.AccountID())
	dsn.APIVersion = c.APIVersion()
	dsn.SupportsZeroImpressions = c.SupportsZeroImpressions()
	dsn.SkipColumnHeader = true

	// Credentials.
	dsn.AccessToken = c.tk.AccessToken
	dsn.ClientID = c.tk.ClientID
	dsn.ClientSecret = c.tk.ClientSecret
	dsn.DeveloperToken = c.tk.DeveloperToken
	dsn.RefreshToken = c.tk.RefreshToken

	return driver.NewDsn(
		c.DatabaseDir(), dsn.String(), c.CacheDir(), c.WithCache(),
	).String()
}

// ExecuteStmt returns the statement to execute.
func (c *Context) ExecuteStmt() string {
	return *c.opts.Query
}

// HistoryFile returns the path to the history file.
func (c *Context) HistoryFile() string {
	if c.homeDir == "" {
		return ""
	}
	return filepath.Join(c.homeDir, "history")
}

// Init retrieves and saves the default authenticate information.
func (c *Context) Init() error {
	// Checks for required flags.
	if err := c.opts.Check(); err != nil {
		return err
	}
	// Checks if it's a supported Adwords API versions.
	// Opens a connection to Awql DB without load anything.
	if _, err := db.Open(c.APIVersion() + ":true|" + c.DatabaseDir()); err != nil {
		return err
	}
	// Create the tool home directory.
	if err := c.mkDirHome(); err != nil {
		return err
	}
	// Checks credential to authenticate to Adwords.
	switch {
	case *c.opts.AccessToken != "":
		// Use temporary credential.
		c.tk = NewTmpCredentials(*c.opts.AccessToken, *c.opts.DeveloperToken)
	case !c.useDefaultAuth():
		// Aks authenticate credentials to user.
		c.tk = AskCredentials()
		if err := c.tk.Save(c.filePath()); err != nil {
			return err
		}
	default:
		// Get authenticate credentials from configuration file.
		c.tk = NewCredentials()
		if err := c.tk.Get(c.filePath()); err != nil {
			return err
		}
	}
	return nil
}

// IsInteractive returns true if query as passed as flag.
func (c *Context) IsInteractive() bool {
	return *c.opts.Query == ""
}

// SupportsZeroImpressions returns true if the support of zero impressions is enable.
func (c *Context) SupportsZeroImpressions() bool {
	return *c.opts.ZeroImpressions
}

// UseBatchMode returns true if raw mode is required.
// It will print results using colon as the column separator, with each row on a new line.
func (c *Context) UseBatchMode() bool {
	return *c.opts.Batch
}

// UseVerboseMode returns true if more output about what the program does is required.
func (c *Context) UseVerboseMode() bool {
	return *c.opts.Verbose
}

// WithAutoRehash returns true if automatic rehashing is enable.
func (c *Context) WithAutoRehash() bool {
	return !*c.opts.NoRehash
}

// WithCache returns true if the cache is enable.
func (c *Context) WithCache() bool {
	return *c.opts.Caching
}

// filePath returns the path to the config file.
func (c *Context) filePath() string {
	if c.homeDir == "" {
		return ""
	}
	return filepath.Join(c.homeDir, "config")
}

// mkDirHome creates if not already exists the home directory.
func (c *Context) mkDirHome() error {
	if c.homeDir != "" {
		// Directory already made.
		return nil
	}
	// Gets user properties.
	usr, err := user.Current()
	if err != nil {
		return err
	}
	// If not exists, create the user home directory of Awql tool: .awql
	c.homeDir = filepath.Join(usr.HomeDir, ".awql")
	if _, err := os.Stat(c.homeDir); os.IsNotExist(err) {
		if err := os.Mkdir(c.homeDir, os.ModePerm); err != nil {
			return err
		}
	}
	// Tries to create the cache directory: .awql/cache
	if err := os.Mkdir(c.CacheDir(), os.ModePerm); !os.IsExist(err) {
		return err
	}
	return nil
}

// useDefaultAuth returns true if the config file exists.
func (c *Context) useDefaultAuth() bool {
	if c.homeDir == "" {
		return false
	}
	// Checks if config file exists.
	_, err := os.Stat(c.filePath())

	return err == nil
}
<commit_msg>Changes the path to the DB schema<commit_after>package conf

import (
	"os"
	"os/user"
	"path/filepath"

	db "github.com/rvflash/awql-db"
	awql "github.com/rvflash/awql-driver"
	"github.com/rvflash/awql/driver"
)

// Options represents all available parameters.
type Options interface {
	AccountID() string
	APIVersion() string
	ExecuteStmt() string
	IsInteractive() bool
	SupportsZeroImpressions() bool
	UseBatchMode() bool
	UseVerboseMode() bool
	WithAutoRehash() bool
	WithCache() bool
}

// Settings represents all configurable options.
type Settings interface {
	Options
	Dsn() string
	CacheDir() string
	DatabaseDir() string
	HistoryFile() string
	Init() error
}

// Context represents the program properties.
type Context struct {
	wrkDir, homeDir string
	tk              *Credentials
	opts            *Flag
}

// New returns an instance of Credentials.
// The string dir contains the path of the current workspace.
func New(dir string) *Context {
	return &Context{wrkDir: dir, opts: Parse()}
}

// AccountID returns the Account ID.
func (c *Context) AccountID() string {
	return *c.opts.AccountID
}

// APIVersion returns the API version.
func (c *Context) APIVersion() string {
	return *c.opts.APIVersion
}

// CacheDir returns the path to store cache files.
func (c *Context) CacheDir() string {
	if c.homeDir == "" {
		return ""
	}
	return filepath.Join(c.homeDir, "cache")
}

// DatabaseDir returns the path to the database.
func (c *Context) DatabaseDir() string {
	return filepath.Join(c.wrkDir, "vendor/github.com/rvflash/awql-db/internal/schema/src")
}

// Dsn outputs the data source name.
func (c *Context) Dsn() string {
	// Data source name used to connect to Adwords.
	dsn := awql.NewDsn(c.AccountID())
	dsn.APIVersion = c.APIVersion()
	dsn.SupportsZeroImpressions = c.SupportsZeroImpressions()
	dsn.SkipColumnHeader = true

	// Credentials.
	dsn.AccessToken = c.tk.AccessToken
	dsn.ClientID = c.tk.ClientID
	dsn.ClientSecret = c.tk.ClientSecret
	dsn.DeveloperToken = c.tk.DeveloperToken
	dsn.RefreshToken = c.tk.RefreshToken

	return driver.NewDsn(
		c.DatabaseDir(), dsn.String(), c.CacheDir(), c.WithCache(),
	).String()
}

// ExecuteStmt returns the statement to execute.
func (c *Context) ExecuteStmt() string {
	return *c.opts.Query
}

// HistoryFile returns the path to the history file.
func (c *Context) HistoryFile() string {
	if c.homeDir == "" {
		return ""
	}
	return filepath.Join(c.homeDir, "history")
}

// Init retrieves and saves the default authenticate information.
func (c *Context) Init() error {
	// Checks for required flags.
	if err := c.opts.Check(); err != nil {
		return err
	}
	// Checks if it's a supported Adwords API versions.
	// Opens a connection to Awql DB without load anything.
	if _, err := db.Open(c.APIVersion() + ":true|" + c.DatabaseDir()); err != nil {
		return err
	}
	// Create the tool home directory.
	if err := c.mkDirHome(); err != nil {
		return err
	}
	// Checks credential to authenticate to Adwords.
	switch {
	case *c.opts.AccessToken != "":
		// Use temporary credential.
		c.tk = NewTmpCredentials(*c.opts.AccessToken, *c.opts.DeveloperToken)
	case !c.useDefaultAuth():
		// Aks authenticate credentials to user.
		c.tk = AskCredentials()
		if err := c.tk.Save(c.filePath()); err != nil {
			return err
		}
	default:
		// Get authenticate credentials from configuration file.
		c.tk = NewCredentials()
		if err := c.tk.Get(c.filePath()); err != nil {
			return err
		}
	}
	return nil
}

// IsInteractive returns true if query as passed as flag.
func (c *Context) IsInteractive() bool {
	return *c.opts.Query == ""
}

// SupportsZeroImpressions returns true if the support of zero impressions is enable.
func (c *Context) SupportsZeroImpressions() bool {
	return *c.opts.ZeroImpressions
}

// UseBatchMode returns true if raw mode is required.
// It will print results using colon as the column separator, with each row on a new line.
func (c *Context) UseBatchMode() bool {
	return *c.opts.Batch
}

// UseVerboseMode returns true if more output about what the program does is required.
func (c *Context) UseVerboseMode() bool {
	return *c.opts.Verbose
}

// WithAutoRehash returns true if automatic rehashing is enable.
func (c *Context) WithAutoRehash() bool {
	return !*c.opts.NoRehash
}

// WithCache returns true if the cache is enable.
func (c *Context) WithCache() bool {
	return *c.opts.Caching
}

// filePath returns the path to the config file.
func (c *Context) filePath() string {
	if c.homeDir == "" {
		return ""
	}
	return filepath.Join(c.homeDir, "config")
}

// mkDirHome creates if not already exists the home directory.
func (c *Context) mkDirHome() error {
	if c.homeDir != "" {
		// Directory already made.
		return nil
	}
	// Gets user properties.
	usr, err := user.Current()
	if err != nil {
		return err
	}
	// If not exists, create the user home directory of Awql tool: .awql
	c.homeDir = filepath.Join(usr.HomeDir, ".awql")
	if _, err := os.Stat(c.homeDir); os.IsNotExist(err) {
		if err := os.Mkdir(c.homeDir, os.ModePerm); err != nil {
			return err
		}
	}
	// Tries to create the cache directory: .awql/cache
	if err := os.Mkdir(c.CacheDir(), os.ModePerm); !os.IsExist(err) {
		return err
	}
	return nil
}

// useDefaultAuth returns true if the config file exists.
func (c *Context) useDefaultAuth() bool {
	if c.homeDir == "" {
		return false
	}
	// Checks if config file exists.
	_, err := os.Stat(c.filePath())

	return err == nil
}
<|endoftext|>


<commit_before>package ssh

const (
	// This is a constant representing a script to install and uninstall public
	// key in remote hosts.
	DefaultPublicKeyInstallScript = `
#!/bin/bash
#
# This is a default script which installs or uninstalls an RSA public key to/from
# authoried_keys file in a typical linux machine.
#
# If the platform differs or if the binaries used in this script are not available
# in targer machine, use the 'install_script' parameter with 'roles/' endpoint to
# register a custom script (applicable for Dynamic type only).
#
# Vault server runs this script on the target machine with the following params:
#
# $1:INSTALL_OPTION: "install" or "uninstall"
#
# $2:PUBLIC_KEY_FILE: File name containing public key to be installed. Vault server
# uses UUID as name to avoid collisions with public keys generated for other requests.
#
# $3:AUTH_KEYS_FILE: Absolute path of the authorized_keys file.
# Currently, vault uses /home/<username>/.ssh/authorized_keys as the path.
#
# [Note: This script will be run by Vault using the registered admin username.
# Notice that some commands below are run as 'sudo'. For graceful execution of
# this script there should not be any password prompts. So, disable password
# prompt for the admin username registered with Vault.

set -e

# Storing arguments into variables, to increase readability of the script.
INSTALL_OPTION=$1
PUBLIC_KEY_FILE=$2
AUTH_KEYS_FILE=$3

# Delete the public key file and the temporary file
function cleanup
{
	rm -f "$PUBLIC_KEY_FILE" temp_$PUBLIC_KEY_FILE
}

# 'cleanup' will be called if the script ends or if any command fails.
trap cleanup EXIT

# Return if the option is anything other than 'install' or 'uninstall'.
if [ "$INSTALL_OPTION" != "install" && "$INSTALL_OPTION" != "uninstall" ]; then
	exit 1
fi

# Remove the key from authorized_keys file if it is already present.
# This step is common for both install and uninstall.
grep -vFf "$PUBLIC_KEY_FILE" "$AUTH_KEYS_FILE" > temp_$PUBLIC_KEY_FILE
cat temp_$PUBLIC_KEY_FILE | sudo tee "$AUTH_KEYS_FILE"

# Append the new public key to authorized_keys file
if [ "$INSTALL_OPTION" == "install" ]; then
	cat "$PUBLIC_KEY_FILE" | sudo tee --append "$AUTH_KEYS_FILE"
fi
`
)
<commit_msg>Vault SSH: Create .ssh directory if not present. Closes #573<commit_after>package ssh

const (
	// This is a constant representing a script to install and uninstall public
	// key in remote hosts.
	DefaultPublicKeyInstallScript = `
#!/bin/bash
#
# This is a default script which installs or uninstalls an RSA public key to/from
# authoried_keys file in a typical linux machine.
#
# If the platform differs or if the binaries used in this script are not available
# in targer machine, use the 'install_script' parameter with 'roles/' endpoint to
# register a custom script (applicable for Dynamic type only).
#
# Vault server runs this script on the target machine with the following params:
#
# $1:INSTALL_OPTION: "install" or "uninstall"
#
# $2:PUBLIC_KEY_FILE: File name containing public key to be installed. Vault server
# uses UUID as name to avoid collisions with public keys generated for other requests.
#
# $3:AUTH_KEYS_FILE: Absolute path of the authorized_keys file.
# Currently, vault uses /home/<username>/.ssh/authorized_keys as the path.
#
# [Note: This script will be run by Vault using the registered admin username.
# Notice that some commands below are run as 'sudo'. For graceful execution of
# this script there should not be any password prompts. So, disable password
# prompt for the admin username registered with Vault.

set -e

# Storing arguments into variables, to increase readability of the script.
INSTALL_OPTION=$1
PUBLIC_KEY_FILE=$2
AUTH_KEYS_FILE=$3

# Delete the public key file and the temporary file
function cleanup
{
	rm -f "$PUBLIC_KEY_FILE" temp_$PUBLIC_KEY_FILE
}

# 'cleanup' will be called if the script ends or if any command fails.
trap cleanup EXIT

# Return if the option is anything other than 'install' or 'uninstall'.
if [ "$INSTALL_OPTION" != "install" ] && [ "$INSTALL_OPTION" != "uninstall" ]; then
	exit 1
fi

# Create the .ssh directory and authorized_keys file if it does not exist
SSH_DIR=$(dirname $AUTH_KEYS_FILE)
sudo mkdir -p "$SSH_DIR"
sudo touch "$AUTH_KEYS_FILE"

# Remove the key from authorized_keys file if it is already present.
# This step is common for both install and uninstall.
grep -vFf "$PUBLIC_KEY_FILE" "$AUTH_KEYS_FILE" > temp_$PUBLIC_KEY_FILE || true
cat temp_$PUBLIC_KEY_FILE | sudo tee "$AUTH_KEYS_FILE"

# Append the new public key to authorized_keys file
if [ "$INSTALL_OPTION" == "install" ]; then
	cat "$PUBLIC_KEY_FILE" | sudo tee --append "$AUTH_KEYS_FILE"
fi
`
)
<|endoftext|>


<commit_before>package common

import (
	"io"
	"strconv"
)

type SerializerOpenTSDB struct {
}

func NewSerializerOpenTSDB() *SerializerOpenTSDB {
	return &SerializerOpenTSDB{}
}

// SerializeOpenTSDBBulk writes Point data to the given writer, conforming to
// the OpenTSDB bulk load protocol (the /api/put endpoint). Note that no line
// has a trailing comma. Downstream programs are responsible for creating
// batches for POSTing using a JSON array, and for adding any trailing commas
// (to conform to JSON). We use only millisecond-precision timestamps.
//
// N.B. OpenTSDB only supports millisecond or second resolution timestamps.
// N.B. OpenTSDB millisecond timestamps must be 13 digits long.
// N.B. OpenTSDB only supports floating-point field values.
//
// This function writes JSON lines that looks like:
// { <metric>, <timestamp>, <value>, <tags> }
//
// For example:
// { "metric": "cpu.usage_user", "timestamp": 14516064000000, "value": 99.5170917755353770, "tags": { "hostname": "host_01", "region": "ap-southeast-2", "datacenter": "ap-southeast-2a" } }
func (m *SerializerOpenTSDB) SerializePoint(w io.Writer, p *Point) (err error) {
	for i := 0; i < len(p.FieldKeys); i++ {
		var value float64
		switch x := p.FieldValues[i].(type) {
		case int:
			value = float64(x)
		case int64:
			value = float64(x)
		case float32:
			value = float64(x)
		case float64:
			value = x
		default:
			panic("bad numeric value for OpenTSDB serialization")
		}

		buf := scratchBufPool.Get().([]byte)
		buf = append(buf, []byte(`{"`)...)
		buf = append(buf, p.MeasurementName...)
		buf = append(buf, '.')
		buf = append(buf, p.FieldKeys[i]...)
		buf = append(buf, []byte(`","timestamp":`)...)
		buf = strconv.AppendInt(buf, p.Timestamp.UTC().UnixNano()/1e6, 10)
		buf = append(buf, []byte(`,"value":`)...)
		buf = strconv.AppendFloat(buf, value, 'f', 16, 64)
		buf = append(buf, []byte(`,"tags":{`)...)
		for i := 0; i < len(p.TagKeys); i++ {
			buf = append(buf, []byte(`"`)...)
			buf = append(buf, p.TagKeys[i]...)
			buf = append(buf, []byte(`":"`)...)
			buf = append(buf, p.TagValues[i]...)
			buf = append(buf, []byte(`"`)...)
			if i+1 != len(p.TagValues) {
				buf = append(buf, ',')
			}
		}
		buf = append(buf, "}}\n"...)
		_, err = w.Write(buf)

		buf = buf[:0]
		scratchBufPool.Put(buf)
	}

	return nil
}

func (s *SerializerOpenTSDB) SerializeSize(w io.Writer, points int64, values int64) error {
	//return serializeSizeInText(w, points, values)
	return nil
}
<commit_msg>fix point protocol<commit_after>package common

import (
	"io"
	"strconv"
)

type SerializerOpenTSDB struct {
}

func NewSerializerOpenTSDB() *SerializerOpenTSDB {
	return &SerializerOpenTSDB{}
}

// SerializeOpenTSDBBulk writes Point data to the given writer, conforming to
// the OpenTSDB bulk load protocol (the /api/put endpoint). Note that no line
// has a trailing comma. Downstream programs are responsible for creating
// batches for POSTing using a JSON array, and for adding any trailing commas
// (to conform to JSON). We use only millisecond-precision timestamps.
//
// N.B. OpenTSDB only supports millisecond or second resolution timestamps.
// N.B. OpenTSDB millisecond timestamps must be 13 digits long.
// N.B. OpenTSDB only supports floating-point field values.
//
// This function writes JSON lines that looks like:
// { <metric>, <timestamp>, <value>, <tags> }
//
// For example:
// { "metric": "cpu.usage_user", "timestamp": 14516064000000, "value": 99.5170917755353770, "tags": { "hostname": "host_01", "region": "ap-southeast-2", "datacenter": "ap-southeast-2a" } }
func (m *SerializerOpenTSDB) SerializePoint(w io.Writer, p *Point) (err error) {
	for i := 0; i < len(p.FieldKeys); i++ {
		var value float64
		switch x := p.FieldValues[i].(type) {
		case int:
			value = float64(x)
		case int64:
			value = float64(x)
		case float32:
			value = float64(x)
		case float64:
			value = x
		default:
			panic("bad numeric value for OpenTSDB serialization")
		}

		buf := scratchBufPool.Get().([]byte)
		buf = append(buf, []byte(`{"metric":"`)...)
		buf = append(buf, p.MeasurementName...)
		buf = append(buf, '.')
		buf = append(buf, p.FieldKeys[i]...)
		buf = append(buf, []byte(`","timestamp":`)...)
		buf = strconv.AppendInt(buf, p.Timestamp.UTC().UnixNano()/1e6, 10)
		buf = append(buf, []byte(`,"value":`)...)
		buf = strconv.AppendFloat(buf, value, 'f', 16, 64)
		buf = append(buf, []byte(`,"tags":{`)...)
		for i := 0; i < len(p.TagKeys); i++ {
			buf = append(buf, []byte(`"`)...)
			buf = append(buf, p.TagKeys[i]...)
			buf = append(buf, []byte(`":"`)...)
			buf = append(buf, p.TagValues[i]...)
			buf = append(buf, []byte(`"`)...)
			if i+1 != len(p.TagValues) {
				buf = append(buf, ',')
			}
		}
		buf = append(buf, "}}\n"...)
		_, err = w.Write(buf)

		buf = buf[:0]
		scratchBufPool.Put(buf)
	}

	return nil
}

func (s *SerializerOpenTSDB) SerializeSize(w io.Writer, points int64, values int64) error {
	//return serializeSizeInText(w, points, values)
	return nil
}
<|endoftext|>


<commit_before>package pserver

type partitioner struct {
	shardNum int
}

// partitioner partitions the parameters into shards.
func newPartitioner(shardNum int) *partitioner {
	return &partitioner{shardNum: shardNum}
}
<commit_msg>remove unused file<commit_after><|endoftext|>


<commit_before>package aio

import (
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/gobottest"
)

var _ gobot.Driver = (*GroveTemperatureSensorDriver)(nil)

func TestGroveTemperatureSensorDriver(t *testing.T) {
	testAdaptor := newAioTestAdaptor()
	d := NewGroveTemperatureSensorDriver(testAdaptor, "123")
	gobottest.Assert(t, d.Connection(), testAdaptor)
	gobottest.Assert(t, d.Pin(), "123")
	gobottest.Assert(t, d.interval, 10*time.Millisecond)
}

func TestGroveTempSensorPublishesTemperatureInCelsius(t *testing.T) {
	sem := make(chan bool, 1)
	a := newAioTestAdaptor()
	d := NewGroveTemperatureSensorDriver(a, "1")

	a.TestAdaptorAnalogRead(func() (val int, err error) {
		val = 585
		return
	})
	gobottest.Assert(t, d.Start(), nil)

	d.Once(d.Event(Data), func(data interface{}) {
		gobottest.Assert(t, fmt.Sprintf("%.2f", data.(float64)), "31.62")
		sem <- true
	})

	select {
	case <-sem:
	case <-time.After(10 * time.Second):
		t.Errorf("Grove Temperature Sensor Event \"Data\" was not published")
	}

	gobottest.Assert(t, d.Temperature(), 31.61532462352477)
}

func TestGroveTempSensorPublishesError(t *testing.T) {
	sem := make(chan bool, 1)
	a := newAioTestAdaptor()
	d := NewGroveTemperatureSensorDriver(a, "1")

	// send error
	a.TestAdaptorAnalogRead(func() (val int, err error) {
		err = errors.New("read error")
		return
	})

	gobottest.Assert(t, d.Start(), nil)

	// expect error
	d.Once(d.Event(Error), func(data interface{}) {
		gobottest.Assert(t, data.(error).Error(), "read error")
		sem <- true
	})

	select {
	case <-sem:
	case <-time.After(time.Second):
		t.Errorf("Grove Temperature Sensor Event \"Error\" was not published")
	}
}

func TestGroveTempSensorHalt(t *testing.T) {
	d := NewGroveTemperatureSensorDriver(newAioTestAdaptor(), "1")
	done := make(chan struct{})
	go func() {
		<-d.halt
		close(done)
	}()
	gobottest.Assert(t, d.Halt(), nil)
	select {
	case <-done:
	case <-time.After(time.Millisecond):
		t.Errorf("Grove Temperature Sensor was not halted")
	}
}

func TestGroveTempDriverDefaultName(t *testing.T) {
	d := NewGroveTemperatureSensorDriver(newAioTestAdaptor(), "1")
	gobottest.Assert(t, strings.HasPrefix(d.Name(), "GroveTemperatureSensor"), true)
}

func TestGroveTempDriverSetName(t *testing.T) {
	d := NewGroveTemperatureSensorDriver(newAioTestAdaptor(), "1")
	d.SetName("mybot")
	gobottest.Assert(t, d.Name(), "mybot")
}
<commit_msg>aio: ensure event handler is setup before starting test<commit_after>package aio

import (
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/gobottest"
)

var _ gobot.Driver = (*GroveTemperatureSensorDriver)(nil)

func TestGroveTemperatureSensorDriver(t *testing.T) {
	testAdaptor := newAioTestAdaptor()
	d := NewGroveTemperatureSensorDriver(testAdaptor, "123")
	gobottest.Assert(t, d.Connection(), testAdaptor)
	gobottest.Assert(t, d.Pin(), "123")
	gobottest.Assert(t, d.interval, 10*time.Millisecond)
}

func TestGroveTempSensorPublishesTemperatureInCelsius(t *testing.T) {
	sem := make(chan bool, 1)
	a := newAioTestAdaptor()
	d := NewGroveTemperatureSensorDriver(a, "1")

	a.TestAdaptorAnalogRead(func() (val int, err error) {
		val = 585
		return
	})
	d.Once(d.Event(Data), func(data interface{}) {
		gobottest.Assert(t, fmt.Sprintf("%.2f", data.(float64)), "31.62")
		sem <- true
	})
	gobottest.Assert(t, d.Start(), nil)

	select {
	case <-sem:
	case <-time.After(10 * time.Second):
		t.Errorf("Grove Temperature Sensor Event \"Data\" was not published")
	}

	gobottest.Assert(t, d.Temperature(), 31.61532462352477)
}

func TestGroveTempSensorPublishesError(t *testing.T) {
	sem := make(chan bool, 1)
	a := newAioTestAdaptor()
	d := NewGroveTemperatureSensorDriver(a, "1")

	// send error
	a.TestAdaptorAnalogRead(func() (val int, err error) {
		err = errors.New("read error")
		return
	})

	gobottest.Assert(t, d.Start(), nil)

	// expect error
	d.Once(d.Event(Error), func(data interface{}) {
		gobottest.Assert(t, data.(error).Error(), "read error")
		sem <- true
	})

	select {
	case <-sem:
	case <-time.After(time.Second):
		t.Errorf("Grove Temperature Sensor Event \"Error\" was not published")
	}
}

func TestGroveTempSensorHalt(t *testing.T) {
	d := NewGroveTemperatureSensorDriver(newAioTestAdaptor(), "1")
	done := make(chan struct{})
	go func() {
		<-d.halt
		close(done)
	}()
	gobottest.Assert(t, d.Halt(), nil)
	select {
	case <-done:
	case <-time.After(time.Millisecond):
		t.Errorf("Grove Temperature Sensor was not halted")
	}
}

func TestGroveTempDriverDefaultName(t *testing.T) {
	d := NewGroveTemperatureSensorDriver(newAioTestAdaptor(), "1")
	gobottest.Assert(t, strings.HasPrefix(d.Name(), "GroveTemperatureSensor"), true)
}

func TestGroveTempDriverSetName(t *testing.T) {
	d := NewGroveTemperatureSensorDriver(newAioTestAdaptor(), "1")
	d.SetName("mybot")
	gobottest.Assert(t, d.Name(), "mybot")
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SurgeMQ Authors. 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 service

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

var (
	bufcnt            int64
	DefaultBufferSize int64

	DeviceInBufferSize  int64
	DeviceOutBufferSize int64

	MasterInBufferSize  int64
	MasterOutBufferSize int64
)

type sequence struct {
	// The current position of the producer or consumer
	cursor,

	// The previous known position of the consumer (if producer) or producer (if consumer)
	gate,

	// These are fillers to pad the cache line, which is generally 64 bytes
	p2, p3, p4, p5, p6, p7 int64
}

func newSequence() *sequence {
	return &sequence{}
}

func (this *sequence) get() int64 {
	return atomic.LoadInt64(&this.cursor)
}

func (this *sequence) set(seq int64) {
	atomic.StoreInt64(&this.cursor, seq)
}

type buffer struct {
	id int64

	readIndex  int64 //读序号
	writeIndex int64 //写序号
	buf        []*[]byte

	size int64
	mask int64

	done int64

	pcond *sync.Cond
	ccond *sync.Cond

	b []byte //readfrom中的临时数组，为反复使用不必创建新数组
}

func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}

	if size == 0 {
		size = DefaultBufferSize
	}

	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		id:         atomic.AddInt64(&bufcnt, 1),
		readIndex:  int64(0),
		writeIndex: int64(0),
		buf:        make([]*[]byte, size),
		size:       size,
		mask:       size - 1,
		pcond:      sync.NewCond(new(sync.Mutex)),
		ccond:      sync.NewCond(new(sync.Mutex)),
		b:          make([]byte, 5),
	}, nil
}

/**
获取当前读序号
*/
func (this *buffer) GetCurrentReadIndex() int64 {
	return atomic.LoadInt64(&this.readIndex)
}

/**
获取当前写序号
*/
func (this *buffer) GetCurrentWriteIndex() int64 {
	return atomic.LoadInt64(&this.writeIndex)
}

func (this *buffer) ID() int64 {
	return this.id
}

func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)

	this.pcond.L.Lock()
	this.ccond.Signal()
	//this.ccond.Broadcast()
	this.pcond.L.Unlock()

	this.ccond.L.Lock()
	this.pcond.Signal()
	//this.pcond.Broadcast()
	this.ccond.L.Unlock()

	return nil
}

/**
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
*/
func (this *buffer) ReadBuffer() (p *[]byte, ok bool) {
	/*this.ccond.L.Lock()
	defer func() {
		//this.pcond.Signal()
		//this.pcond.Broadcast()
		this.ccond.L.Unlock()
		//time.Sleep(3 * time.Millisecond)
	}()*/
	defer func() {
		this.ccond.L.Lock()
		this.pcond.Signal()
		this.ccond.L.Unlock()
	}()
	ok = false
	p = nil
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return nil, false
		}
		writeIndex = this.GetCurrentWriteIndex()
		if readIndex >= writeIndex {
			this.ccond.L.Lock()
			this.pcond.Signal()
			//this.pcond.Broadcast()
			this.ccond.Wait()
			this.ccond.L.Unlock()
			//runtime.Gosched()
			//time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
		//time.Sleep(500 * time.Microsecond)
		//time.Sleep(1 * time.Millisecond)
	}
	index := readIndex & this.mask //替代求模

	p = this.buf[index]
	this.buf[index] = nil

	atomic.AddInt64(&this.readIndex, int64(1))
	if p != nil {
		ok = true
	}
	return p, ok
}

/**
写入ringbuffer指针，以及将写序号加1
*/
func (this *buffer) WriteBuffer(in *[]byte) (ok bool) {
	/*this.pcond.L.Lock()
	defer func() {
		//this.ccond.Signal()
		//this.ccond.Broadcast()
		this.pcond.L.Unlock()
		//time.Sleep(3 * time.Millisecond)
	}()*/
	defer func() {
		this.pcond.L.Lock()
		this.ccond.Signal()
		this.pcond.L.Unlock()
	}()
	ok = false
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return false
		}
		readIndex = this.GetCurrentReadIndex()
		if writeIndex >= readIndex && writeIndex-readIndex >= this.size {
			this.pcond.L.Lock()
			this.ccond.Signal()
			//this.ccond.Broadcast()
			this.pcond.Wait()
			this.pcond.L.Unlock()
			//time.Sleep(1 * time.Millisecond)
			//runtime.Gosched()
			//time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
		//time.Sleep(500 * time.Microsecond)
		//time.Sleep(1 * time.Millisecond)
	}
	index := writeIndex & this.mask //替代求模

	this.buf[index] = in

	atomic.AddInt64(&this.writeIndex, int64(1))
	ok = true
	return ok
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		time.Sleep(10 * time.Millisecond)
		if this.isDone() {
			return total, io.EOF
		}

		var write_bytes []byte

		//b := make([]byte, 5)
		n, err := r.Read(this.b[0:1])
		if err != nil {
			return total, io.EOF
		}
		total += int64(n)
		max_cnt := 1
		for {
			if this.isDone() {
				return total, io.EOF
			}
			// If we have read 5 bytes and still not done, then there's a problem.
			if max_cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}
			_, err := r.Read(this.b[max_cnt:(max_cnt + 1)])

			//fmt.Println(b)
			if err != nil {
				return total, err
			}
			if this.b[max_cnt] >= 0x80 {
				max_cnt++
			} else {
				break
			}
		}
		remlen, m := binary.Uvarint(this.b[1 : max_cnt+1])
		remlen_tmp := int64(remlen)
		start_ := int64(1) + int64(m)
		total_tmp := remlen_tmp + start_

		write_bytes = make([]byte, total_tmp)
		copy(write_bytes[0:m+1], this.b[0:m+1])
		nlen := int64(0)
		times := 0
		cnt_ := int64(32)
		for nlen < remlen_tmp {
			if this.isDone() {
				return total, io.EOF
			}
			if times > 100 {
				return total, io.EOF
			} else {
				times = 0
			}
			times++
			tmpm := remlen_tmp - nlen

			b_ := write_bytes[(start_ + nlen):]
			if tmpm > cnt_ {
				b_ = write_bytes[(start_ + nlen):(start_ + nlen + cnt_)]
			}

			//b_ := make([]byte, remlen)
			n, err = r.Read(b_[0:])

			if err != nil {
				/*Log.Errorc(func() string {
					return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
				})
				time.Sleep(5 * time.Millisecond)
				continue*/
				return total, err
			}
			//write_bytes = append(write_bytes, b_[0:]...)
			nlen += int64(n)
			total += int64(n)
		}

		ok := this.WriteBuffer(&write_bytes)

		if !ok {
			return total, errors.New("write ringbuffer failed")
		}
	}
}

func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		if this.isDone() {
			return total, io.EOF
		}

		p, ok := this.ReadBuffer()
		if !ok {
			return total, errors.New("read buffer failed")
		}
		// There's some data, let's process it first
		if len(*p) > 0 {
			n, err := w.Write(*p)
			total += int64(n)
			//Log.Debugc(func() string{ return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)})

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

	}
}

func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func ringCopy(dst, src []byte, start int64) int {
	n := len(src)

	i, l := 0, 0

	for n > 0 {
		l = copy(dst[start:], src[i:])
		i += l
		n -= l

		if n > 0 {
			start = 0
		}
	}

	return i
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n&(n-1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom_not_receiver(r io.Reader) (*[]byte, error) {

	total := int64(0)

	if this.isDone() {
		return nil, io.EOF
	}

	var write_bytes []byte

	//b := make([]byte, 5)
	n, err := r.Read(this.b[0:1])
	if err != nil {
		return total, io.EOF
	}
	total += int64(n)
	max_cnt := 1
	for {
		if this.isDone() {
			return nil, io.EOF
		}
		// If we have read 5 bytes and still not done, then there's a problem.
		if max_cnt > 4 {
			return nil, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
		}
		_, err := r.Read(this.b[max_cnt:(max_cnt + 1)])

		//fmt.Println(b)
		if err != nil {
			return nil, err
		}
		if this.b[max_cnt] >= 0x80 {
			max_cnt++
		} else {
			break
		}
	}
	remlen, m := binary.Uvarint(this.b[1 : max_cnt+1])
	remlen_tmp := int64(remlen)
	start_ := int64(1) + int64(m)
	total_tmp := remlen_tmp + start_

	write_bytes = make([]byte, total_tmp)
	copy(write_bytes[0:m+1], this.b[0:m+1])
	nlen := int64(0)
	times := 0
	cnt_ := int64(32)
	for nlen < remlen_tmp {
		if this.isDone() {
			return nil, io.EOF
		}
		if times > 100 {
			return nil, io.EOF
		} else {
			times = 0
		}
		times++
		tmpm := remlen_tmp - nlen

		b_ := write_bytes[(start_ + nlen):]
		if tmpm > cnt_ {
			b_ = write_bytes[(start_ + nlen):(start_ + nlen + cnt_)]
		}

		//b_ := make([]byte, remlen)
		n, err = r.Read(b_[0:])

		if err != nil {
			/*Log.Errorc(func() string {
				return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
			})
			time.Sleep(5 * time.Millisecond)
			continue*/
			return nil, err
		}
		//write_bytes = append(write_bytes, b_[0:]...)
		nlen += int64(n)
		total += int64(n)
	}

	//ok := this.WriteBuffer(&write_bytes)

	return &write_bytes, nil
}
<commit_msg>修改 buffer.go process.go sendrecv.go xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 去掉receiver<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. 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 service

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

var (
	bufcnt            int64
	DefaultBufferSize int64

	DeviceInBufferSize  int64
	DeviceOutBufferSize int64

	MasterInBufferSize  int64
	MasterOutBufferSize int64
)

type sequence struct {
	// The current position of the producer or consumer
	cursor,

	// The previous known position of the consumer (if producer) or producer (if consumer)
	gate,

	// These are fillers to pad the cache line, which is generally 64 bytes
	p2, p3, p4, p5, p6, p7 int64
}

func newSequence() *sequence {
	return &sequence{}
}

func (this *sequence) get() int64 {
	return atomic.LoadInt64(&this.cursor)
}

func (this *sequence) set(seq int64) {
	atomic.StoreInt64(&this.cursor, seq)
}

type buffer struct {
	id int64

	readIndex  int64 //读序号
	writeIndex int64 //写序号
	buf        []*[]byte

	size int64
	mask int64

	done int64

	pcond *sync.Cond
	ccond *sync.Cond

	b []byte //readfrom中的临时数组，为反复使用不必创建新数组
}

func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}

	if size == 0 {
		size = DefaultBufferSize
	}

	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		id:         atomic.AddInt64(&bufcnt, 1),
		readIndex:  int64(0),
		writeIndex: int64(0),
		buf:        make([]*[]byte, size),
		size:       size,
		mask:       size - 1,
		pcond:      sync.NewCond(new(sync.Mutex)),
		ccond:      sync.NewCond(new(sync.Mutex)),
		b:          make([]byte, 5),
	}, nil
}

/**
获取当前读序号
*/
func (this *buffer) GetCurrentReadIndex() int64 {
	return atomic.LoadInt64(&this.readIndex)
}

/**
获取当前写序号
*/
func (this *buffer) GetCurrentWriteIndex() int64 {
	return atomic.LoadInt64(&this.writeIndex)
}

func (this *buffer) ID() int64 {
	return this.id
}

func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)

	this.pcond.L.Lock()
	this.ccond.Signal()
	//this.ccond.Broadcast()
	this.pcond.L.Unlock()

	this.ccond.L.Lock()
	this.pcond.Signal()
	//this.pcond.Broadcast()
	this.ccond.L.Unlock()

	return nil
}

/**
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
*/
func (this *buffer) ReadBuffer() (p *[]byte, ok bool) {
	/*this.ccond.L.Lock()
	defer func() {
		//this.pcond.Signal()
		//this.pcond.Broadcast()
		this.ccond.L.Unlock()
		//time.Sleep(3 * time.Millisecond)
	}()*/
	defer func() {
		this.ccond.L.Lock()
		this.pcond.Signal()
		this.ccond.L.Unlock()
	}()
	ok = false
	p = nil
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return nil, false
		}
		writeIndex = this.GetCurrentWriteIndex()
		if readIndex >= writeIndex {
			this.ccond.L.Lock()
			this.pcond.Signal()
			//this.pcond.Broadcast()
			this.ccond.Wait()
			this.ccond.L.Unlock()
			//runtime.Gosched()
			//time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
		//time.Sleep(500 * time.Microsecond)
		//time.Sleep(1 * time.Millisecond)
	}
	index := readIndex & this.mask //替代求模

	p = this.buf[index]
	this.buf[index] = nil

	atomic.AddInt64(&this.readIndex, int64(1))
	if p != nil {
		ok = true
	}
	return p, ok
}

/**
写入ringbuffer指针，以及将写序号加1
*/
func (this *buffer) WriteBuffer(in *[]byte) (ok bool) {
	/*this.pcond.L.Lock()
	defer func() {
		//this.ccond.Signal()
		//this.ccond.Broadcast()
		this.pcond.L.Unlock()
		//time.Sleep(3 * time.Millisecond)
	}()*/
	defer func() {
		this.pcond.L.Lock()
		this.ccond.Signal()
		this.pcond.L.Unlock()
	}()
	ok = false
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return false
		}
		readIndex = this.GetCurrentReadIndex()
		if writeIndex >= readIndex && writeIndex-readIndex >= this.size {
			this.pcond.L.Lock()
			this.ccond.Signal()
			//this.ccond.Broadcast()
			this.pcond.Wait()
			this.pcond.L.Unlock()
			//time.Sleep(1 * time.Millisecond)
			//runtime.Gosched()
			//time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
		//time.Sleep(500 * time.Microsecond)
		//time.Sleep(1 * time.Millisecond)
	}
	index := writeIndex & this.mask //替代求模

	this.buf[index] = in

	atomic.AddInt64(&this.writeIndex, int64(1))
	ok = true
	return ok
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		time.Sleep(10 * time.Millisecond)
		if this.isDone() {
			return total, io.EOF
		}

		var write_bytes []byte

		//b := make([]byte, 5)
		n, err := r.Read(this.b[0:1])
		if err != nil {
			return total, io.EOF
		}
		total += int64(n)
		max_cnt := 1
		for {
			if this.isDone() {
				return total, io.EOF
			}
			// If we have read 5 bytes and still not done, then there's a problem.
			if max_cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}
			_, err := r.Read(this.b[max_cnt:(max_cnt + 1)])

			//fmt.Println(b)
			if err != nil {
				return total, err
			}
			if this.b[max_cnt] >= 0x80 {
				max_cnt++
			} else {
				break
			}
		}
		remlen, m := binary.Uvarint(this.b[1 : max_cnt+1])
		remlen_tmp := int64(remlen)
		start_ := int64(1) + int64(m)
		total_tmp := remlen_tmp + start_

		write_bytes = make([]byte, total_tmp)
		copy(write_bytes[0:m+1], this.b[0:m+1])
		nlen := int64(0)
		times := 0
		cnt_ := int64(32)
		for nlen < remlen_tmp {
			if this.isDone() {
				return total, io.EOF
			}
			if times > 100 {
				return total, io.EOF
			} else {
				times = 0
			}
			times++
			tmpm := remlen_tmp - nlen

			b_ := write_bytes[(start_ + nlen):]
			if tmpm > cnt_ {
				b_ = write_bytes[(start_ + nlen):(start_ + nlen + cnt_)]
			}

			//b_ := make([]byte, remlen)
			n, err = r.Read(b_[0:])

			if err != nil {
				/*Log.Errorc(func() string {
					return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
				})
				time.Sleep(5 * time.Millisecond)
				continue*/
				return total, err
			}
			//write_bytes = append(write_bytes, b_[0:]...)
			nlen += int64(n)
			total += int64(n)
		}

		ok := this.WriteBuffer(&write_bytes)

		if !ok {
			return total, errors.New("write ringbuffer failed")
		}
	}
}

func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		if this.isDone() {
			return total, io.EOF
		}

		p, ok := this.ReadBuffer()
		if !ok {
			return total, errors.New("read buffer failed")
		}
		// There's some data, let's process it first
		if len(*p) > 0 {
			n, err := w.Write(*p)
			total += int64(n)
			//Log.Debugc(func() string{ return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)})

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

	}
}

func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func ringCopy(dst, src []byte, start int64) int {
	n := len(src)

	i, l := 0, 0

	for n > 0 {
		l = copy(dst[start:], src[i:])
		i += l
		n -= l

		if n > 0 {
			start = 0
		}
	}

	return i
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n&(n-1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom_not_receiver(r io.Reader) (*[]byte, error) {

	total := int64(0)

	if this.isDone() {
		return nil, io.EOF
	}

	var write_bytes []byte

	//b := make([]byte, 5)
	n, err := r.Read(this.b[0:1])
	if err != nil {
		return nil, io.EOF
	}
	total += int64(n)
	max_cnt := 1
	for {
		if this.isDone() {
			return nil, io.EOF
		}
		// If we have read 5 bytes and still not done, then there's a problem.
		if max_cnt > 4 {
			return nil, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
		}
		_, err := r.Read(this.b[max_cnt:(max_cnt + 1)])

		//fmt.Println(b)
		if err != nil {
			return nil, err
		}
		if this.b[max_cnt] >= 0x80 {
			max_cnt++
		} else {
			break
		}
	}
	remlen, m := binary.Uvarint(this.b[1 : max_cnt+1])
	remlen_tmp := int64(remlen)
	start_ := int64(1) + int64(m)
	total_tmp := remlen_tmp + start_

	write_bytes = make([]byte, total_tmp)
	copy(write_bytes[0:m+1], this.b[0:m+1])
	nlen := int64(0)
	times := 0
	cnt_ := int64(32)
	for nlen < remlen_tmp {
		if this.isDone() {
			return nil, io.EOF
		}
		if times > 100 {
			return nil, io.EOF
		} else {
			times = 0
		}
		times++
		tmpm := remlen_tmp - nlen

		b_ := write_bytes[(start_ + nlen):]
		if tmpm > cnt_ {
			b_ = write_bytes[(start_ + nlen):(start_ + nlen + cnt_)]
		}

		//b_ := make([]byte, remlen)
		n, err = r.Read(b_[0:])

		if err != nil {
			/*Log.Errorc(func() string {
				return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
			})
			time.Sleep(5 * time.Millisecond)
			continue*/
			return nil, err
		}
		//write_bytes = append(write_bytes, b_[0:]...)
		nlen += int64(n)
		total += int64(n)
	}

	//ok := this.WriteBuffer(&write_bytes)

	return &write_bytes, nil
}
<|endoftext|>


<commit_before>// +build !jenkins

package content

import (
	"fmt"
	"os"
	"testing"

	"github.com/Financial-Times/neo-utils-go/neoutils"
	"github.com/jmcvetta/neoism"
	"github.com/stretchr/testify/assert"
)

const (
	contentUUID 		= "ce3f2f5e-33d1-4c36-89e3-51aa00fd5660"
	conceptUUID             = "412e4ca3-f8d5-4456-8606-064c1dba3c45"
	noBodyContentUuid	= "6440aa4a-1298-4a49-9346-78d546bc0229"
)

var contentWithoutABody = content{
	UUID:  noBodyContentUuid,
	Title: "Missing Body",
}

var standardContent = content{
	UUID:          contentUUID,
	Title:         "Content Title",
	PublishedDate: "1970-01-01T01:00:00.000Z",
	Body:          "Some body",
}

var shorterContent = content{
	UUID: contentUUID,
	Body: "With No Publish Date and No Title",
}

var updatedContent = content{
	UUID: contentUUID,
	Title: "New Ttitle",
	PublishedDate: "1999-12-12T01:00:00.000Z",
	Body: "Doesn't matter",
}

func TestDeleteWithNoRelsIsDeleted(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")

	deleted, err := contentDriver.Delete(standardContent.UUID)
	assert.True(deleted, "Didn't manage to delete content for uuid %s", standardContent.UUID)
	assert.NoError(err, "Error deleting content for uuid %s", standardContent.UUID)

	c, deleted, err := contentDriver.Read(standardContent.UUID)

	assert.Equal(content{}, c, "Found content %s who should have been deleted", c)
	assert.False(deleted, "Found content for uuid %s who should have been deleted", standardContent.UUID)
	assert.NoError(err, "Error trying to find content for uuid %s", standardContent.UUID)
}

func TestDeleteWithRelsBecomesThing(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")
	writeRelationship(db, standardContent.UUID, conceptUUID, t, assert)

	deleted, err := contentDriver.Delete(standardContent.UUID)
	assert.NoError(err, "Error deleting content for uuid %s", standardContent.UUID)
	assert.True(deleted, "Didn't manage to delete content for uuid %s", standardContent.UUID)

	c, found, err := contentDriver.Read(standardContent.UUID)

	assert.Equal(content{}, c, "Found content %s who should have been deleted", c)
	assert.False(found, "Found content for uuid %s who should have been deleted", standardContent.UUID)
	assert.NoError(err, "Error trying to find content for uuid %s", standardContent.UUID)

	exists, err := doesThingExist(standardContent.UUID, db)
	assert.NoError(err)
	assert.True(exists, "Failed to find Thing for deleted content with relationships")
}

func TestCreateAllValuesPresent(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")

	storedContent, _, err := contentDriver.Read(standardContent.UUID)

	assert.NoError(err)
	assert.NotEmpty(storedContent, "Failed to retireve stored content")
	assert.Equal(storedContent.(content).UUID, standardContent.UUID, "Failed to match UUID")
	assert.Equal(storedContent.(content).Title, standardContent.Title, "Failed to match Title")
	assert.Equal(storedContent.(content).PublishedDate, standardContent.PublishedDate, "Failed to match PublishedDate")
}

func TestCreateNotAllValuesPresent(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(shorterContent), "Failed to write content")

	storedContent, _, err := contentDriver.Read(shorterContent.UUID)

	assert.NoError(err)
	assert.Empty(storedContent.(content).PublishedDate)
	assert.Empty(storedContent.(content).Title)

}

func TestCreateDoesNotStoreBody(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")
	storedContent, _, err := contentDriver.Read(standardContent.UUID)

	assert.NoError(err)
	assert.NotEmpty(storedContent, "Failed to retrive stored content")
	assert.Empty(storedContent.(content).Body, "Body should not have been stored")

}

func TestWillUpdateProperties(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")
	storedContent, _, err := contentDriver.Read(contentUUID)

	assert.NoError(err)
	assert.Equal(storedContent.(content).Title, standardContent.Title)
	assert.Equal(storedContent.(content).PublishedDate, standardContent.PublishedDate)

	assert.NoError(contentDriver.Write(updatedContent), "Failed to write updated content")
	storedContent, _, err = contentDriver.Read(contentUUID)

	assert.NoError(err)
	assert.Equal(storedContent.(content).Title, updatedContent.Title, "Should have updated Title but it is still present")
	assert.Equal(storedContent.(content).PublishedDate, updatedContent.PublishedDate, "Should have updated PublishedDate but it is still present")
}

func TestUpdateWillRemovePropertiesNoLongerPresent(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")
	storedContent, _, err := contentDriver.Read(contentUUID)

	assert.NoError(err)
	assert.NotEmpty(storedContent.(content).Title)
	assert.NotEmpty(storedContent.(content).PublishedDate)

	assert.NoError(contentDriver.Write(shorterContent), "Failed to write updated content")
	storedContent, _, err = contentDriver.Read(contentUUID)

	assert.NoError(err)
	assert.NotEmpty(storedContent, "Failed to rtreive updated content")
	assert.Empty(storedContent.(content).Title, "Update should have removed Title but it is still present")
	assert.Empty(storedContent.(content).PublishedDate, "Update should have removed PublishedDate but it is still present")
}

func TestWriteCalculateEpocCorrectly(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	uuid := standardContent.UUID
	contentReceived := content{UUID: uuid, Title: "TestContent", PublishedDate: "1970-01-01T01:00:00.000Z", Body: "Some Test text"}
	contentDriver.Write(contentReceived)

	result := []struct {
		PublishedDateEpoc int `json:"t.publishedDateEpoch"`
	}{}

	getEpochQuery := &neoism.CypherQuery{
		Statement: `
			MATCH (t:Content {uuid:{uuid}}) RETURN t.publishedDateEpoch
			`,
		Parameters: neoism.Props{
			"uuid": standardContent.UUID,
		},
		Result: &result,
	}

	err := contentDriver.conn.CypherBatch([]*neoism.CypherQuery{getEpochQuery})
	assert.NoError(err)
	assert.Equal(3600, result[0].PublishedDateEpoc, "Epoc of 1970-01-01T01:00:00.000Z should be 3600")
}

func TestWritePrefLabelIsAlsoWrittenAndIsEqualToTitle(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	contentDriver.Write(standardContent)

	result := []struct {
		PrefLabel string `json:"t.prefLabel"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Content {uuid:{uuid}}) RETURN t.prefLabel
				`,
		Parameters: neoism.Props{
			"uuid": standardContent.UUID,
		},
		Result: &result,
	}

	err := contentDriver.conn.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal(standardContent.Title, result[0].PrefLabel, "PrefLabel should be 'Content Title'")
}

func TestContentWontBeWrittenIfNoBody(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(contentWithoutABody), "Failed to write content")
	storedContent, _, err := contentDriver.Read(contentWithoutABody.UUID)

	assert.NoError(err)
	assert.Equal(content{}, storedContent, "No content should be written when the content has no body")
}

func getDatabaseConnectionAndCheckClean(t *testing.T, assert *assert.Assertions) neoutils.NeoConnection {
	db := getDatabaseConnection(assert)
	cleanDB(db, t, assert)
	checkDbClean(db, t)
	return db
}

func getDatabaseConnection(assert *assert.Assertions) neoutils.NeoConnection {
	url := os.Getenv("NEO4J_TEST_URL")
	if url == "" {
		url = "http://localhost:7474/db/data"
	}

	conf := neoutils.DefaultConnectionConfig()
	conf.Transactional = false
	db, err := neoutils.Connect(url, conf)
	assert.NoError(err, "Failed to connect to Neo4j")
	return db
}

func cleanDB(db neoutils.CypherRunner, t *testing.T, assert *assert.Assertions) {
	qs := []*neoism.CypherQuery{
		{
			Statement: fmt.Sprintf("MATCH (mc:Thing {uuid: '%v'}) DETACH DELETE mc", conceptUUID),
		},
		{
			Statement: fmt.Sprintf("MATCH (fc:Thing {uuid: '%v'}) DETACH DELETE fc", contentUUID),
		},
	}

	err := db.CypherBatch(qs)
	assert.NoError(err)
}

func writeRelationship(db neoutils.NeoConnection, contentId string, conceptId string, t *testing.T, assert *assert.Assertions) {
	var annotateQuery string
	var qs []*neoism.CypherQuery

	annotateQuery = `
                MERGE (content:Thing{uuid:{contentId}})
                MERGE (concept:Thing) ON CREATE SET concept.uuid = {conceptId}
                MERGE (content)-[pred:SOME_PPREDICATE]->(concept)
          `

	qs = []*neoism.CypherQuery{
		{
			Statement:  annotateQuery,
			Parameters: neoism.Props{"contentId": contentId, "conceptId": conceptId},
		},
	}

	err := db.CypherBatch(qs)
	assert.NoError(err)
}

func doesThingExist(uuid string, db neoutils.NeoConnection) (bool, error) {

	result := []struct{
		UUID  string  `json:"uuid,omitempty"`
	}{}
	query := &neoism.CypherQuery{
		Statement: "MATCH (t:Thing {uuid:{uuid}}) RETURN t.uuid as uuid",
		Parameters: neoism.Props{
			"uuid": uuid,
		},
		Result: &result,
	}
	err := db.CypherBatch([]*neoism.CypherQuery{query})

	return len(result) > 0, err
}

func checkDbClean(db neoutils.CypherRunner, t *testing.T) {
	assert := assert.New(t)

	result := []struct {
		Uuid string `json:"t.uuid"`
	}{}

	checkGraph := neoism.CypherQuery{
		Statement: `
			MATCH (t:Thing) WHERE t.uuid in {uuids} RETURN t.uuid
		`,
		Parameters: neoism.Props{
			"uuids": []string{standardContent.UUID, conceptUUID },
		},
		Result: &result,
	}
	err := db.CypherBatch([]*neoism.CypherQuery{&checkGraph})
	assert.NoError(err)
	assert.Empty(result)
}

func getCypherDriver(db neoutils.NeoConnection) service {
	cr := NewCypherContentService(db)
	cr.Initialise()
	return cr
}
<commit_msg>removed test about not storing Body and moved check to TestAllValuesPresent<commit_after>// +build !jenkins

package content

import (
	"fmt"
	"os"
	"testing"

	"github.com/Financial-Times/neo-utils-go/neoutils"
	"github.com/jmcvetta/neoism"
	"github.com/stretchr/testify/assert"
)

const (
	contentUUID 		= "ce3f2f5e-33d1-4c36-89e3-51aa00fd5660"
	conceptUUID             = "412e4ca3-f8d5-4456-8606-064c1dba3c45"
	noBodyContentUuid	= "6440aa4a-1298-4a49-9346-78d546bc0229"
)

var contentWithoutABody = content{
	UUID:  noBodyContentUuid,
	Title: "Missing Body",
}

var standardContent = content{
	UUID:          contentUUID,
	Title:         "Content Title",
	PublishedDate: "1970-01-01T01:00:00.000Z",
	Body:          "Some body",
}

var shorterContent = content{
	UUID: contentUUID,
	Body: "With No Publish Date and No Title",
}

var updatedContent = content{
	UUID: contentUUID,
	Title: "New Ttitle",
	PublishedDate: "1999-12-12T01:00:00.000Z",
	Body: "Doesn't matter",
}

func TestDeleteWithNoRelsIsDeleted(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")

	deleted, err := contentDriver.Delete(standardContent.UUID)
	assert.True(deleted, "Didn't manage to delete content for uuid %s", standardContent.UUID)
	assert.NoError(err, "Error deleting content for uuid %s", standardContent.UUID)

	c, deleted, err := contentDriver.Read(standardContent.UUID)

	assert.Equal(content{}, c, "Found content %s who should have been deleted", c)
	assert.False(deleted, "Found content for uuid %s who should have been deleted", standardContent.UUID)
	assert.NoError(err, "Error trying to find content for uuid %s", standardContent.UUID)
}

func TestDeleteWithRelsBecomesThing(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")
	writeRelationship(db, standardContent.UUID, conceptUUID, t, assert)

	deleted, err := contentDriver.Delete(standardContent.UUID)
	assert.NoError(err, "Error deleting content for uuid %s", standardContent.UUID)
	assert.True(deleted, "Didn't manage to delete content for uuid %s", standardContent.UUID)

	c, found, err := contentDriver.Read(standardContent.UUID)

	assert.Equal(content{}, c, "Found content %s who should have been deleted", c)
	assert.False(found, "Found content for uuid %s who should have been deleted", standardContent.UUID)
	assert.NoError(err, "Error trying to find content for uuid %s", standardContent.UUID)

	exists, err := doesThingExist(standardContent.UUID, db)
	assert.NoError(err)
	assert.True(exists, "Failed to find Thing for deleted content with relationships")
}

func TestCreateAllValuesPresent(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")

	storedContent, _, err := contentDriver.Read(standardContent.UUID)

	assert.NoError(err)
	assert.NotEmpty(storedContent, "Failed to retireve stored content")
	assert.Equal(storedContent.(content).UUID, standardContent.UUID, "Failed to match UUID")
	assert.Equal(storedContent.(content).Title, standardContent.Title, "Failed to match Title")
	assert.Equal(storedContent.(content).PublishedDate, standardContent.PublishedDate, "Failed to match PublishedDate")
	assert.Empty(storedContent.(content).Body, "Body should not have been stored")

}

func TestCreateNotAllValuesPresent(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(shorterContent), "Failed to write content")

	storedContent, _, err := contentDriver.Read(shorterContent.UUID)

	assert.NoError(err)
	assert.Empty(storedContent.(content).PublishedDate)
	assert.Empty(storedContent.(content).Title)

}

func TestWillUpdateProperties(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")
	storedContent, _, err := contentDriver.Read(contentUUID)

	assert.NoError(err)
	assert.Equal(storedContent.(content).Title, standardContent.Title)
	assert.Equal(storedContent.(content).PublishedDate, standardContent.PublishedDate)

	assert.NoError(contentDriver.Write(updatedContent), "Failed to write updated content")
	storedContent, _, err = contentDriver.Read(contentUUID)

	assert.NoError(err)
	assert.Equal(storedContent.(content).Title, updatedContent.Title, "Should have updated Title but it is still present")
	assert.Equal(storedContent.(content).PublishedDate, updatedContent.PublishedDate, "Should have updated PublishedDate but it is still present")
}

func TestUpdateWillRemovePropertiesNoLongerPresent(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(standardContent), "Failed to write content")
	storedContent, _, err := contentDriver.Read(contentUUID)

	assert.NoError(err)
	assert.NotEmpty(storedContent.(content).Title)
	assert.NotEmpty(storedContent.(content).PublishedDate)

	assert.NoError(contentDriver.Write(shorterContent), "Failed to write updated content")
	storedContent, _, err = contentDriver.Read(contentUUID)

	assert.NoError(err)
	assert.NotEmpty(storedContent, "Failed to rtreive updated content")
	assert.Empty(storedContent.(content).Title, "Update should have removed Title but it is still present")
	assert.Empty(storedContent.(content).PublishedDate, "Update should have removed PublishedDate but it is still present")
}

func TestWriteCalculateEpocCorrectly(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	uuid := standardContent.UUID
	contentReceived := content{UUID: uuid, Title: "TestContent", PublishedDate: "1970-01-01T01:00:00.000Z", Body: "Some Test text"}
	contentDriver.Write(contentReceived)

	result := []struct {
		PublishedDateEpoc int `json:"t.publishedDateEpoch"`
	}{}

	getEpochQuery := &neoism.CypherQuery{
		Statement: `
			MATCH (t:Content {uuid:{uuid}}) RETURN t.publishedDateEpoch
			`,
		Parameters: neoism.Props{
			"uuid": standardContent.UUID,
		},
		Result: &result,
	}

	err := contentDriver.conn.CypherBatch([]*neoism.CypherQuery{getEpochQuery})
	assert.NoError(err)
	assert.Equal(3600, result[0].PublishedDateEpoc, "Epoc of 1970-01-01T01:00:00.000Z should be 3600")
}

func TestWritePrefLabelIsAlsoWrittenAndIsEqualToTitle(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	contentDriver.Write(standardContent)

	result := []struct {
		PrefLabel string `json:"t.prefLabel"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Content {uuid:{uuid}}) RETURN t.prefLabel
				`,
		Parameters: neoism.Props{
			"uuid": standardContent.UUID,
		},
		Result: &result,
	}

	err := contentDriver.conn.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal(standardContent.Title, result[0].PrefLabel, "PrefLabel should be 'Content Title'")
}

func TestContentWontBeWrittenIfNoBody(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	contentDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(contentDriver.Write(contentWithoutABody), "Failed to write content")
	storedContent, _, err := contentDriver.Read(contentWithoutABody.UUID)

	assert.NoError(err)
	assert.Equal(content{}, storedContent, "No content should be written when the content has no body")
}

func getDatabaseConnectionAndCheckClean(t *testing.T, assert *assert.Assertions) neoutils.NeoConnection {
	db := getDatabaseConnection(assert)
	cleanDB(db, t, assert)
	checkDbClean(db, t)
	return db
}

func getDatabaseConnection(assert *assert.Assertions) neoutils.NeoConnection {
	url := os.Getenv("NEO4J_TEST_URL")
	if url == "" {
		url = "http://localhost:7474/db/data"
	}

	conf := neoutils.DefaultConnectionConfig()
	conf.Transactional = false
	db, err := neoutils.Connect(url, conf)
	assert.NoError(err, "Failed to connect to Neo4j")
	return db
}

func cleanDB(db neoutils.CypherRunner, t *testing.T, assert *assert.Assertions) {
	qs := []*neoism.CypherQuery{
		{
			Statement: fmt.Sprintf("MATCH (mc:Thing {uuid: '%v'}) DETACH DELETE mc", conceptUUID),
		},
		{
			Statement: fmt.Sprintf("MATCH (fc:Thing {uuid: '%v'}) DETACH DELETE fc", contentUUID),
		},
	}

	err := db.CypherBatch(qs)
	assert.NoError(err)
}

func writeRelationship(db neoutils.NeoConnection, contentId string, conceptId string, t *testing.T, assert *assert.Assertions) {
	var annotateQuery string
	var qs []*neoism.CypherQuery

	annotateQuery = `
                MERGE (content:Thing{uuid:{contentId}})
                MERGE (concept:Thing) ON CREATE SET concept.uuid = {conceptId}
                MERGE (content)-[pred:SOME_PPREDICATE]->(concept)
          `

	qs = []*neoism.CypherQuery{
		{
			Statement:  annotateQuery,
			Parameters: neoism.Props{"contentId": contentId, "conceptId": conceptId},
		},
	}

	err := db.CypherBatch(qs)
	assert.NoError(err)
}

func doesThingExist(uuid string, db neoutils.NeoConnection) (bool, error) {

	result := []struct{
		UUID  string  `json:"uuid,omitempty"`
	}{}
	query := &neoism.CypherQuery{
		Statement: "MATCH (t:Thing {uuid:{uuid}}) RETURN t.uuid as uuid",
		Parameters: neoism.Props{
			"uuid": uuid,
		},
		Result: &result,
	}
	err := db.CypherBatch([]*neoism.CypherQuery{query})

	return len(result) > 0, err
}

func checkDbClean(db neoutils.CypherRunner, t *testing.T) {
	assert := assert.New(t)

	result := []struct {
		Uuid string `json:"t.uuid"`
	}{}

	checkGraph := neoism.CypherQuery{
		Statement: `
			MATCH (t:Thing) WHERE t.uuid in {uuids} RETURN t.uuid
		`,
		Parameters: neoism.Props{
			"uuids": []string{standardContent.UUID, conceptUUID },
		},
		Result: &result,
	}
	err := db.CypherBatch([]*neoism.CypherQuery{&checkGraph})
	assert.NoError(err)
	assert.Empty(result)
}

func getCypherDriver(db neoutils.NeoConnection) service {
	cr := NewCypherContentService(db)
	cr.Initialise()
	return cr
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Jorge Acereda Maciá. All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE file.

// Package ebml decodes EBML data.
//
// EBML is short for Extensible Binary Meta Language. EBML specifies a
// binary and octet (byte) aligned format inspired by the principle of
// XML. EBML itself is a generalized description of the technique of
// binary markup. Like XML, it is completely agnostic to any data that it
// might contain.
// For a specification, see http://ebml.sourceforge.net/specs/
package ebml

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"reflect"
	"strconv"
)

// ReachedPayloadError is generated when a field tagged with
// ebmlstop:"1" is reached.
type ReachedPayloadError struct {
	First *Element
	Rest  *Element
}

func (r ReachedPayloadError) Error() string {
	return "Reached payload"
}

// Element represents an EBML-encoded chunk of data.
type Element struct {
	R  io.Reader
	Id uint
}

func (e *Element) String() string {
	return fmt.Sprintf("{%+v %x}", e.R, e.Id)
}

// Size returns the size of the element.
func (e *Element) Size() int64 {
	lr := e.R.(*io.LimitedReader)
	return lr.N
}

// Creates the root element corresponding to the data available in r.
func RootElement(r io.Reader) (*Element, error) {
	e := &Element{io.LimitReader(r, math.MaxInt64), 0}
	return e, nil
}

func remaining(x int8) (rem int) {
	for x > 0 {
		rem++
		x += x
	}
	return
}

func readVint(r io.Reader) (val uint64, err error, rem int) {
	v := make([]uint8, 1)
	_, err = io.ReadFull(r, v)
	if err == nil {
		val = uint64(v[0])
		rem = remaining(int8(val))
		for i := 0; err == nil && i < rem; i++ {
			_, err = io.ReadFull(r, v)
			val <<= 8
			val += uint64(v[0])
		}
	}
	return
}

func readSize(r io.Reader) (int64, error) {
	val, err, rem := readVint(r)
	return int64(val & ^(128 << uint(rem*8-rem))), err
}

// Next returns the next child element in an element.
func (e *Element) Next() (*Element, error) {
	var ne Element
	id, err, _ := readVint(e.R)
	if err != nil {
		return nil, err
	}
	var sz int64
	sz, err = readSize(e.R)
	if err != nil {
		return nil, err
	}
	ne.R = io.LimitReader(e.R, sz)
	ne.Id = uint(id)
	return &ne, err
}

func (e *Element) readUint64() (uint64, error) {
	d, err := e.ReadData()
	var i int
	sz := len(d)
	var val uint64
	for i = 0; i < sz; i++ {
		val <<= 8
		val += uint64(d[i])
	}
	return val, err
}

func (e *Element) readUint() (uint, error) {
	val, err := e.readUint64()
	return uint(val), err
}

func (e *Element) readString() (string, error) {
	s, err := e.ReadData()
	return string(s), err
}

func (e *Element) ReadData() (d []byte, err error) {
	sz := e.Size()
	d = make([]uint8, sz, sz)
	_, err = io.ReadFull(e.R, d)
	return
}

func (e *Element) readFloat() (val float64, err error) {
	var uval uint64
	uval, err = e.readUint64()
	if e.Size() == 8 {
		val = math.Float64frombits(uval)
	} else {
		val = float64(math.Float32frombits(uint32(uval)))
	}
	return
}

func (e *Element) skip() (err error) {
	_, err = e.ReadData()
	return
}

// Unmarshal reads EBML data from r into data. Data must be a pointer
// to a struct. Fields present in the struct but absent in the stream
// will just keep their zero value.
// Returns an error that can be an io.Error or a ReachedPayloadError
// containing the first element and the the parent element containing
// the rest of the elements.
func (e *Element) Unmarshal(val interface{}) error {
	return e.readStruct(reflect.Indirect(reflect.ValueOf(val)))
}

func getTag(f reflect.StructField, s string) uint {
	sid := f.Tag.Get(s)
	id, _ := strconv.ParseUint(sid, 16, 0)
	return uint(id)
}

func lookup(reqid uint, t reflect.Type) int {
	for i, l := 0, t.NumField(); i < l; i++ {
		f := t.Field(i)
		if getTag(f, "ebml") == reqid {
			return i - 1000000*int(getTag(f, "ebmlstop"))
		}
	}
	return -1
}

func setDefaults(v reflect.Value) {
	t := v.Type()
	for i, l := 0, t.NumField(); i < l; i++ {
		fv := v.Field(i)
		switch fv.Kind() {
		case reflect.Int:
			fallthrough
		case reflect.Int8:
			fallthrough
		case reflect.Int16:
			fallthrough
		case reflect.Int32:
			fallthrough
		case reflect.Int64:
			fallthrough
		case reflect.Uint:
			fallthrough
		case reflect.Uint8:
			fallthrough
		case reflect.Uint16:
			fallthrough
		case reflect.Uint32:
			fallthrough
		case reflect.Uint64:
			fallthrough
		case reflect.Float32:
			fallthrough
		case reflect.Float64:
			fallthrough
		case reflect.String:
			setFieldDefaults(fv, t.Field(i), v)
		case reflect.Array:
			fallthrough
		case reflect.Struct:
			fallthrough
		case reflect.Slice:
			break
		default:
			log.Panic("Unsupported type")
		}
	}
}

func setFieldDefaults(v reflect.Value, sf reflect.StructField, s reflect.Value) {
	if v.CanInterface() && reflect.DeepEqual(
		v.Interface(), reflect.Zero(v.Type()).Interface()) {
		tag := sf.Tag.Get("ebmldef")
		if tag != "" {
			switch v.Kind() {
			case reflect.Int:
				fallthrough
			case reflect.Int64:
				u, _ := strconv.ParseInt(tag, 10, 0)
				v.SetInt(int64(u))
			case reflect.Uint:
				fallthrough
			case reflect.Uint64:
				u, _ := strconv.ParseUint(tag, 10, 0)
				v.SetUint(u)
			case reflect.Float32:
				fallthrough
			case reflect.Float64:
				f, _ := strconv.ParseFloat(tag, 64)
				v.SetFloat(f)
			case reflect.String:
				v.SetString(tag)
			default:
				log.Panic("Unsupported default value")
			}
		}
		ltag := sf.Tag.Get("ebmldeflink")
		if ltag != "" {
			v.Set(s.FieldByName(ltag))
		}
	}
}

func (e *Element) readStruct(v reflect.Value) (err error) {
	t := v.Type()
	for err == nil {
		var ne *Element
		ne, err = e.Next()
		if err == io.EOF {
			err = nil
			break
		}
		i := lookup(ne.Id, t)
		if i >= 0 {
			err = ne.readField(v.Field(i))
		} else if i == -1 {
			err = ne.skip()
		} else {
			err = ReachedPayloadError{ne, e}
		}
	}
	setDefaults(v)
	return
}

func (e *Element) readField(v reflect.Value) (err error) {
	switch v.Kind() {
	case reflect.Struct:
		err = e.readStruct(v)
	case reflect.Slice:
		err = e.readSlice(v)
	case reflect.Array:
		for i, l := 0, v.Len(); i < l && err == nil; i++ {
			err = e.readStruct(v.Index(i))
		}
	case reflect.String:
		var s string
		s, err = e.readString()
		v.SetString(s)
	case reflect.Int:
		fallthrough
	case reflect.Int64:
		var u uint64
		u, err = e.readUint64()
		v.SetInt(int64(u))
	case reflect.Uint:
		fallthrough
	case reflect.Uint64:
		var u uint64
		u, err = e.readUint64()
		v.SetUint(u)
	case reflect.Float32:
		fallthrough
	case reflect.Float64:
		var f float64
		f, err = e.readFloat()
		v.SetFloat(f)
	default:
		err = errors.New("Unknown type: " + v.String())
	}
	return
}

func (e *Element) readSlice(v reflect.Value) (err error) {
	switch v.Type().Elem().Kind() {
	case reflect.Uint8:
		var sl []uint8
		sl, err = e.ReadData()
		if err == nil {
			v.Set(reflect.ValueOf(sl))
		}
	case reflect.Struct:
		vl := v.Len()
		ne := reflect.New(v.Type().Elem())
		nsl := reflect.Append(v, reflect.Indirect(ne))
		v.Set(nsl)
		err = e.readStruct(v.Index(vl))
	default:
		err = errors.New("Unknown slice type: " + v.String())
	}
	return
}
<commit_msg>Bug reading floats<commit_after>// Copyright (c) 2012, Jorge Acereda Maciá. All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE file.

// Package ebml decodes EBML data.
//
// EBML is short for Extensible Binary Meta Language. EBML specifies a
// binary and octet (byte) aligned format inspired by the principle of
// XML. EBML itself is a generalized description of the technique of
// binary markup. Like XML, it is completely agnostic to any data that it
// might contain.
// For a specification, see http://ebml.sourceforge.net/specs/
package ebml

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"reflect"
	"strconv"
)

// ReachedPayloadError is generated when a field tagged with
// ebmlstop:"1" is reached.
type ReachedPayloadError struct {
	First *Element
	Rest  *Element
}

func (r ReachedPayloadError) Error() string {
	return "Reached payload"
}

// Element represents an EBML-encoded chunk of data.
type Element struct {
	R  io.Reader
	Id uint
}

func (e *Element) String() string {
	return fmt.Sprintf("{%+v %x}", e.R, e.Id)
}

// Size returns the size of the element.
func (e *Element) Size() int64 {
	lr := e.R.(*io.LimitedReader)
	return lr.N
}

// Creates the root element corresponding to the data available in r.
func RootElement(r io.Reader) (*Element, error) {
	e := &Element{io.LimitReader(r, math.MaxInt64), 0}
	return e, nil
}

func remaining(x int8) (rem int) {
	for x > 0 {
		rem++
		x += x
	}
	return
}

func readVint(r io.Reader) (val uint64, err error, rem int) {
	v := make([]uint8, 1)
	_, err = io.ReadFull(r, v)
	if err == nil {
		val = uint64(v[0])
		rem = remaining(int8(val))
		for i := 0; err == nil && i < rem; i++ {
			_, err = io.ReadFull(r, v)
			val <<= 8
			val += uint64(v[0])
		}
	}
	return
}

func readSize(r io.Reader) (int64, error) {
	val, err, rem := readVint(r)
	return int64(val & ^(128 << uint(rem*8-rem))), err
}

// Next returns the next child element in an element.
func (e *Element) Next() (*Element, error) {
	var ne Element
	id, err, _ := readVint(e.R)
	if err != nil {
		return nil, err
	}
	var sz int64
	sz, err = readSize(e.R)
	if err != nil {
		return nil, err
	}
	ne.R = io.LimitReader(e.R, sz)
	ne.Id = uint(id)
	return &ne, err
}

func (e *Element) readUint64() (uint64, error) {
	d, err := e.ReadData()
	var i int
	sz := len(d)
	var val uint64
	for i = 0; i < sz; i++ {
		val <<= 8
		val += uint64(d[i])
	}
	return val, err
}

func (e *Element) readUint() (uint, error) {
	val, err := e.readUint64()
	return uint(val), err
}

func (e *Element) readString() (string, error) {
	s, err := e.ReadData()
	return string(s), err
}

func (e *Element) ReadData() (d []byte, err error) {
	sz := e.Size()
	d = make([]uint8, sz, sz)
	_, err = io.ReadFull(e.R, d)
	return
}

func (e *Element) readFloat() (val float64, err error) {
	var uval uint64
	sz := e.Size()
	uval, err = e.readUint64()
	if sz == 8 {
		val = math.Float64frombits(uval)
	} else {
		val = float64(math.Float32frombits(uint32(uval)))
	}
	return
}

func (e *Element) skip() (err error) {
	_, err = e.ReadData()
	return
}

// Unmarshal reads EBML data from r into data. Data must be a pointer
// to a struct. Fields present in the struct but absent in the stream
// will just keep their zero value.
// Returns an error that can be an io.Error or a ReachedPayloadError
// containing the first element and the the parent element containing
// the rest of the elements.
func (e *Element) Unmarshal(val interface{}) error {
	return e.readStruct(reflect.Indirect(reflect.ValueOf(val)))
}

func getTag(f reflect.StructField, s string) uint {
	sid := f.Tag.Get(s)
	id, _ := strconv.ParseUint(sid, 16, 0)
	return uint(id)
}

func lookup(reqid uint, t reflect.Type) int {
	for i, l := 0, t.NumField(); i < l; i++ {
		f := t.Field(i)
		if getTag(f, "ebml") == reqid {
			return i - 1000000*int(getTag(f, "ebmlstop"))
		}
	}
	return -1
}

func setDefaults(v reflect.Value) {
	t := v.Type()
	for i, l := 0, t.NumField(); i < l; i++ {
		fv := v.Field(i)
		switch fv.Kind() {
		case reflect.Int:
			fallthrough
		case reflect.Int8:
			fallthrough
		case reflect.Int16:
			fallthrough
		case reflect.Int32:
			fallthrough
		case reflect.Int64:
			fallthrough
		case reflect.Uint:
			fallthrough
		case reflect.Uint8:
			fallthrough
		case reflect.Uint16:
			fallthrough
		case reflect.Uint32:
			fallthrough
		case reflect.Uint64:
			fallthrough
		case reflect.Float32:
			fallthrough
		case reflect.Float64:
			fallthrough
		case reflect.String:
			setFieldDefaults(fv, t.Field(i), v)
		case reflect.Array:
			fallthrough
		case reflect.Struct:
			fallthrough
		case reflect.Slice:
			break
		default:
			log.Panic("Unsupported type")
		}
	}
}

func setFieldDefaults(v reflect.Value, sf reflect.StructField, s reflect.Value) {
	if v.CanInterface() && reflect.DeepEqual(
		v.Interface(), reflect.Zero(v.Type()).Interface()) {
		tag := sf.Tag.Get("ebmldef")
		if tag != "" {
			switch v.Kind() {
			case reflect.Int:
				fallthrough
			case reflect.Int64:
				u, _ := strconv.ParseInt(tag, 10, 0)
				v.SetInt(int64(u))
			case reflect.Uint:
				fallthrough
			case reflect.Uint64:
				u, _ := strconv.ParseUint(tag, 10, 0)
				v.SetUint(u)
			case reflect.Float32:
				fallthrough
			case reflect.Float64:
				f, _ := strconv.ParseFloat(tag, 64)
				v.SetFloat(f)
			case reflect.String:
				v.SetString(tag)
			default:
				log.Panic("Unsupported default value")
			}
		}
		ltag := sf.Tag.Get("ebmldeflink")
		if ltag != "" {
			v.Set(s.FieldByName(ltag))
		}
	}
}

func (e *Element) readStruct(v reflect.Value) (err error) {
	t := v.Type()
	for err == nil {
		var ne *Element
		ne, err = e.Next()
		if err == io.EOF {
			err = nil
			break
		}
		i := lookup(ne.Id, t)
		if i >= 0 {
			err = ne.readField(v.Field(i))
		} else if i == -1 {
			err = ne.skip()
		} else {
			err = ReachedPayloadError{ne, e}
		}
	}
	setDefaults(v)
	return
}

func (e *Element) readField(v reflect.Value) (err error) {
	switch v.Kind() {
	case reflect.Struct:
		err = e.readStruct(v)
	case reflect.Slice:
		err = e.readSlice(v)
	case reflect.Array:
		for i, l := 0, v.Len(); i < l && err == nil; i++ {
			err = e.readStruct(v.Index(i))
		}
	case reflect.String:
		var s string
		s, err = e.readString()
		v.SetString(s)
	case reflect.Int:
		fallthrough
	case reflect.Int64:
		var u uint64
		u, err = e.readUint64()
		v.SetInt(int64(u))
	case reflect.Uint:
		fallthrough
	case reflect.Uint64:
		var u uint64
		u, err = e.readUint64()
		v.SetUint(u)
	case reflect.Float32:
		fallthrough
	case reflect.Float64:
		var f float64
		f, err = e.readFloat()
		v.SetFloat(f)
	default:
		err = errors.New("Unknown type: " + v.String())
	}
	return
}

func (e *Element) readSlice(v reflect.Value) (err error) {
	switch v.Type().Elem().Kind() {
	case reflect.Uint8:
		var sl []uint8
		sl, err = e.ReadData()
		if err == nil {
			v.Set(reflect.ValueOf(sl))
		}
	case reflect.Struct:
		vl := v.Len()
		ne := reflect.New(v.Type().Elem())
		nsl := reflect.Append(v, reflect.Indirect(ne))
		v.Set(nsl)
		err = e.readStruct(v.Index(vl))
	default:
		err = errors.New("Unknown slice type: " + v.String())
	}
	return
}
<|endoftext|>