<commit_before>package builder

import (
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"

	"github.com/spf13/afero"
	"github.com/spf13/fsync"

	"github.com/aymerick/kowa/models"
	"github.com/aymerick/kowa/utils"
)

const (
	IMAGES_DIR    = "img"
	THEMES_DIR    = "themes"
	TEMPLATES_DIR = "templates"
	PARTIALS_DIR  = "partials"
	ASSETS_DIR    = "assets"

	DEFAULT_THEME = "test"

	MAX_SLUG        = 50
	MAX_PAST_EVENTS = 5
)

var registeredNodeBuilders = make(map[string]func(*SiteBuilder) NodeBuilder)

// Registers a node builder
func RegisterNodeBuilder(name string, initializer func(*SiteBuilder) NodeBuilder) {
	if _, exists := registeredNodeBuilders[name]; exists {
		panic(fmt.Sprintf("Builder initializer already registered: %s", name))
	}

	registeredNodeBuilders[name] = initializer
}

type SiteBuilder struct {
	site   *models.Site
	config *SiteBuilderConfig

	images         []*models.Image
	errorCollector *ErrorCollector

	// cache for #layout method
	masterLayout *template.Template

	// internal vars
	nodeBuilders map[string]NodeBuilder

	siteVars *SiteVars
	tplDir   string
}

type SiteBuilderConfig struct {
	WorkingDir string
	OutputDir  string
	Theme      string
	UglyURL    bool
}

func NewSiteBuilder(site *models.Site, config *SiteBuilderConfig) *SiteBuilder {
	result := &SiteBuilder{
		site:   site,
		config: config,

		errorCollector: NewErrorCollector(),

		nodeBuilders: make(map[string]NodeBuilder),
	}

	result.initBuilders()
	result.setTemplatesDir()

	return result
}

// Theme used by builder
func (builder *SiteBuilder) Theme() string {
	result := builder.config.Theme
	if result == "" {
		result = builder.site.Theme
	}

	if result == "" {
		result = DEFAULT_THEME
	}

	return result
}

// Are we building site with ugly urls ?
func (builder *SiteBuilder) UglyUrl() bool {
	if builder.config.UglyURL {
		return true
	} else {
		return builder.site.UglyURL
	}
}

// Build site
func (builder *SiteBuilder) Build() {
	// load nodes
	if builder.loadNodes(); builder.HaveError() {
		return
	}

	// compute site variables
	if builder.fillSiteVars(); builder.HaveError() {
		return
	}

	// generate nodes
	if builder.generateNodes(); builder.HaveError() {
		return
	}

	// copy images
	if builder.copyCollectedImages(); builder.HaveError() {
		return
	}

	// copy assets
	builder.copyAssets()
}

// Initialize builders
func (builder *SiteBuilder) initBuilders() {
	for name, initializer := range registeredNodeBuilders {
		builder.nodeBuilders[name] = initializer(builder)
	}
}

// Find templates dir
func (builder *SiteBuilder) setTemplatesDir() {
	dirPath := path.Join(builder.themeDir(), TEMPLATES_DIR)
	if _, err := os.Stat(dirPath); os.IsNotExist(err) {
		// no /templates subdir found
		builder.tplDir = builder.themeDir()
	} else {
		builder.tplDir = dirPath
	}
}

// Get given node builder
func (builder *SiteBuilder) nodeBuilder(name string) NodeBuilder {
	return builder.nodeBuilders[name]
}

// Load nodes
func (builder *SiteBuilder) loadNodes() {
	for _, nodeBuilder := range builder.nodeBuilders {
		nodeBuilder.Load()
	}
}

// Returns nodes to display in navigation bar
func (builder *SiteBuilder) navBarNodes() []*Node {
	result := []*Node{}

	for _, nodeBuilder := range builder.nodeBuilders {
		nodes := nodeBuilder.NavBarNodes()
		if len(nodes) > 0 {
			result = append(result, nodes...)
		}
	}

	return result
}

// Returns all activities vars (to display in template)
func (builder *SiteBuilder) activitiesVars() []*ActivityVars {
	nodeBuilder := builder.nodeBuilder(KIND_ACTIVITIES)

	activities, ok := nodeBuilder.Data("activities").([]*ActivityVars)
	if !ok {
		panic("This should never happen")
	}

	return activities
}

// Fill site variables
func (builder *SiteBuilder) fillSiteVars() {
	builder.siteVars = NewSiteVars(builder)
	builder.siteVars.fill()
}

// Generate nodes
func (builder *SiteBuilder) generateNodes() {
	for _, nodeBuilder := range builder.nodeBuilders {
		nodeBuilder.Generate()
	}
}

// Copy images
func (builder *SiteBuilder) copyCollectedImages() {
	errStep := "Copy images"

	imgDir := builder.genImagesDir()

	// ensure img dir
	if err := builder.ensureDir(imgDir); err != nil {
		builder.addError(errStep, err)
		return
	}

	// copy images to img dir
	for _, image := range builder.images {
		for _, derivative := range models.Derivatives {
			srcFile := image.DerivativeFilePath(derivative)
			if err := builder.copyFile(srcFile, imgDir); err != nil {
				builder.addError(errStep, err)
			}
		}
	}
}

// Copy theme assets
func (builder *SiteBuilder) copyAssets() error {
	syncer := fsync.NewSyncer()
	syncer.SrcFs = new(afero.OsFs)
	syncer.DestFs = new(afero.OsFs)

	return syncer.Sync(builder.genAssetsDir(), builder.themeAssetsDir())
}

// Collect image, and returns the template vars for that image
func (builder *SiteBuilder) addImage(img *models.Image) *ImageVars {
	builder.images = append(builder.images, img)

	return NewImageVars(img)
}

// Check if builder have error
func (builder *SiteBuilder) HaveError() bool {
	return builder.errorCollector.ErrorsNb > 0
}

// Add an error to collector
func (builder *SiteBuilder) addError(step string, err error) {
	builder.errorCollector.addError(step, err)
}

// Add an error when generating a node
func (builder *SiteBuilder) addNodeBuilderError(nodeKind string, err error) {
	step := fmt.Sprintf("Building %s", nodeKind)
	builder.addError(step, err)
}

// Dump errors
func (builder *SiteBuilder) DumpErrors() {
	builder.errorCollector.dump()
}

// Computes theme directory
func (builder *SiteBuilder) themeDir() string {
	return path.Join(builder.config.WorkingDir, THEMES_DIR, builder.Theme())
}

// Computes theme templates directory
func (builder *SiteBuilder) themeTemplatesDir() string {
	if builder.tplDir == "" {
		panic("Templates directory not set")
	}

	return builder.tplDir
}

// Computes theme assets directory
func (builder *SiteBuilder) themeAssetsDir() string {
	return path.Join(builder.themeDir(), ASSETS_DIR)
}

// Computes directory where site is generated
func (builder *SiteBuilder) GenDir() string {
	return path.Join(builder.config.WorkingDir, builder.config.OutputDir)
}

// Computes directory where images are copied
func (builder *SiteBuilder) genImagesDir() string {
	return path.Join(builder.GenDir(), IMAGES_DIR)
}

// Computes directory where assets are copied
func (builder *SiteBuilder) genAssetsDir() string {
	return path.Join(builder.GenDir(), ASSETS_DIR)
}

// Compute template path for given template name
func (builder *SiteBuilder) templatePath(tplName string) string {
	return path.Join(builder.themeTemplatesDir(), fmt.Sprintf("%s.html", tplName))
}

// Returns partials directory path
func (builder *SiteBuilder) partialsPath() string {
	return path.Join(builder.themeTemplatesDir(), PARTIALS_DIR)
}

func (builder *SiteBuilder) setupLayout() *template.Template {
	errStep := "template init"

	// parse layout
	result, err := template.ParseFiles(builder.templatePath("layout"))
	if err != nil {
		builder.addError(errStep, err)
		return nil
	}

	// setup FuncMap
	result.Funcs(builder.FuncMap())

	// load partials
	partialDir := builder.partialsPath()

	files, err := ioutil.ReadDir(partialDir)
	if err != nil && err != os.ErrExist {
		builder.addError(errStep, err)
	} else {
		for _, file := range files {
			fileName := file.Name()

			if !file.IsDir() && strings.HasSuffix(fileName, ".html") {
				filePath := path.Join(partialDir, fileName)

				// read partial
				binData, err := ioutil.ReadFile(filePath)
				if err != nil {
					builder.addError(errStep, err)
				} else {
					// eg: partials/navbar
					tplName := fmt.Sprintf("%s/%s", PARTIALS_DIR, utils.FileBase(fileName))

					// add partial to layout
					_, err := result.New(tplName).Parse(string(binData))
					if err != nil {
						builder.addError(errStep, err)
					}
				}
			}
		}
	}

	return result
}

// Get master layout template
func (builder *SiteBuilder) layout() *template.Template {
	if builder.masterLayout == nil {
		builder.masterLayout = builder.setupLayout()
	}

	return builder.masterLayout
}

// Dump templates
func (builder *SiteBuilder) DumpLayout() {
	log.Printf("Layout templates:")
	for _, tpl := range builder.masterLayout.Templates() {
		log.Printf("  -> %s", tpl.Name())
	}
}

// Prune directories for given absolute dir path
func (builder *SiteBuilder) ensureDir(dirPath string) error {
	// log.Printf("[DBG] Creating dir: %s", dirPath)

	err := os.MkdirAll(dirPath, 0777)
	if err != nil && err != os.ErrExist {
		return err
	}

	return err
}

// Prune directories for given absolute file path
func (builder *SiteBuilder) ensureFileDir(osPath string) error {
	return builder.ensureDir(path.Dir(osPath))
}

// Computes local file path for given URL
func (builder *SiteBuilder) filePath(fullUrl string) string {
	return path.Join(builder.GenDir(), fullUrl)
}

// Copy file to given directory
func (builder *SiteBuilder) copyFile(fromFilePath string, toDir string) error {
	// open source file
	src, err := os.Open(fromFilePath)
	if err != nil {
		return err
	}
	defer src.Close()

	// open destination file
	dstFilePath := path.Join(toDir, path.Base(fromFilePath))

	dst, err := os.Create(dstFilePath)
	if err != nil {
		return err
	}
	defer dst.Close()

	// copy
	if _, err := io.Copy(dst, src); err != nil {
		return err
	}

	return nil
}
<commit_msg>Fixes default theme<commit_after>package builder

import (
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"

	"github.com/spf13/afero"
	"github.com/spf13/fsync"

	"github.com/aymerick/kowa/models"
	"github.com/aymerick/kowa/utils"
)

const (
	IMAGES_DIR    = "img"
	THEMES_DIR    = "themes"
	TEMPLATES_DIR = "templates"
	PARTIALS_DIR  = "partials"
	ASSETS_DIR    = "assets"

	DEFAULT_THEME = "willy" // @todo FIXME !

	MAX_SLUG        = 50
	MAX_PAST_EVENTS = 5
)

var registeredNodeBuilders = make(map[string]func(*SiteBuilder) NodeBuilder)

// Registers a node builder
func RegisterNodeBuilder(name string, initializer func(*SiteBuilder) NodeBuilder) {
	if _, exists := registeredNodeBuilders[name]; exists {
		panic(fmt.Sprintf("Builder initializer already registered: %s", name))
	}

	registeredNodeBuilders[name] = initializer
}

type SiteBuilder struct {
	site   *models.Site
	config *SiteBuilderConfig

	images         []*models.Image
	errorCollector *ErrorCollector

	// cache for #layout method
	masterLayout *template.Template

	// internal vars
	nodeBuilders map[string]NodeBuilder

	siteVars *SiteVars
	tplDir   string
}

type SiteBuilderConfig struct {
	WorkingDir string
	OutputDir  string
	Theme      string
	UglyURL    bool
}

func NewSiteBuilder(site *models.Site, config *SiteBuilderConfig) *SiteBuilder {
	result := &SiteBuilder{
		site:   site,
		config: config,

		errorCollector: NewErrorCollector(),

		nodeBuilders: make(map[string]NodeBuilder),
	}

	result.initBuilders()
	result.setTemplatesDir()

	return result
}

// Theme used by builder
func (builder *SiteBuilder) Theme() string {
	result := builder.config.Theme
	if result == "" {
		result = builder.site.Theme
	}

	if result == "" {
		result = DEFAULT_THEME
	}

	return result
}

// Are we building site with ugly urls ?
func (builder *SiteBuilder) UglyUrl() bool {
	if builder.config.UglyURL {
		return true
	} else {
		return builder.site.UglyURL
	}
}

// Build site
func (builder *SiteBuilder) Build() {
	// load nodes
	if builder.loadNodes(); builder.HaveError() {
		return
	}

	// compute site variables
	if builder.fillSiteVars(); builder.HaveError() {
		return
	}

	// generate nodes
	if builder.generateNodes(); builder.HaveError() {
		return
	}

	// copy images
	if builder.copyCollectedImages(); builder.HaveError() {
		return
	}

	// copy assets
	builder.copyAssets()
}

// Initialize builders
func (builder *SiteBuilder) initBuilders() {
	for name, initializer := range registeredNodeBuilders {
		builder.nodeBuilders[name] = initializer(builder)
	}
}

// Find templates dir
func (builder *SiteBuilder) setTemplatesDir() {
	dirPath := path.Join(builder.themeDir(), TEMPLATES_DIR)
	if _, err := os.Stat(dirPath); os.IsNotExist(err) {
		// no /templates subdir found
		builder.tplDir = builder.themeDir()
	} else {
		builder.tplDir = dirPath
	}
}

// Get given node builder
func (builder *SiteBuilder) nodeBuilder(name string) NodeBuilder {
	return builder.nodeBuilders[name]
}

// Load nodes
func (builder *SiteBuilder) loadNodes() {
	for _, nodeBuilder := range builder.nodeBuilders {
		nodeBuilder.Load()
	}
}

// Returns nodes to display in navigation bar
func (builder *SiteBuilder) navBarNodes() []*Node {
	result := []*Node{}

	for _, nodeBuilder := range builder.nodeBuilders {
		nodes := nodeBuilder.NavBarNodes()
		if len(nodes) > 0 {
			result = append(result, nodes...)
		}
	}

	return result
}

// Returns all activities vars (to display in template)
func (builder *SiteBuilder) activitiesVars() []*ActivityVars {
	nodeBuilder := builder.nodeBuilder(KIND_ACTIVITIES)

	activities, ok := nodeBuilder.Data("activities").([]*ActivityVars)
	if !ok {
		panic("This should never happen")
	}

	return activities
}

// Fill site variables
func (builder *SiteBuilder) fillSiteVars() {
	builder.siteVars = NewSiteVars(builder)
	builder.siteVars.fill()
}

// Generate nodes
func (builder *SiteBuilder) generateNodes() {
	for _, nodeBuilder := range builder.nodeBuilders {
		nodeBuilder.Generate()
	}
}

// Copy images
func (builder *SiteBuilder) copyCollectedImages() {
	errStep := "Copy images"

	imgDir := builder.genImagesDir()

	// ensure img dir
	if err := builder.ensureDir(imgDir); err != nil {
		builder.addError(errStep, err)
		return
	}

	// copy images to img dir
	for _, image := range builder.images {
		for _, derivative := range models.Derivatives {
			srcFile := image.DerivativeFilePath(derivative)
			if err := builder.copyFile(srcFile, imgDir); err != nil {
				builder.addError(errStep, err)
			}
		}
	}
}

// Copy theme assets
func (builder *SiteBuilder) copyAssets() error {
	syncer := fsync.NewSyncer()
	syncer.SrcFs = new(afero.OsFs)
	syncer.DestFs = new(afero.OsFs)

	return syncer.Sync(builder.genAssetsDir(), builder.themeAssetsDir())
}

// Collect image, and returns the template vars for that image
func (builder *SiteBuilder) addImage(img *models.Image) *ImageVars {
	builder.images = append(builder.images, img)

	return NewImageVars(img)
}

// Check if builder have error
func (builder *SiteBuilder) HaveError() bool {
	return builder.errorCollector.ErrorsNb > 0
}

// Add an error to collector
func (builder *SiteBuilder) addError(step string, err error) {
	builder.errorCollector.addError(step, err)
}

// Add an error when generating a node
func (builder *SiteBuilder) addNodeBuilderError(nodeKind string, err error) {
	step := fmt.Sprintf("Building %s", nodeKind)
	builder.addError(step, err)
}

// Dump errors
func (builder *SiteBuilder) DumpErrors() {
	builder.errorCollector.dump()
}

// Computes theme directory
func (builder *SiteBuilder) themeDir() string {
	return path.Join(builder.config.WorkingDir, THEMES_DIR, builder.Theme())
}

// Computes theme templates directory
func (builder *SiteBuilder) themeTemplatesDir() string {
	if builder.tplDir == "" {
		panic("Templates directory not set")
	}

	return builder.tplDir
}

// Computes theme assets directory
func (builder *SiteBuilder) themeAssetsDir() string {
	return path.Join(builder.themeDir(), ASSETS_DIR)
}

// Computes directory where site is generated
func (builder *SiteBuilder) GenDir() string {
	return path.Join(builder.config.WorkingDir, builder.config.OutputDir)
}

// Computes directory where images are copied
func (builder *SiteBuilder) genImagesDir() string {
	return path.Join(builder.GenDir(), IMAGES_DIR)
}

// Computes directory where assets are copied
func (builder *SiteBuilder) genAssetsDir() string {
	return path.Join(builder.GenDir(), ASSETS_DIR)
}

// Compute template path for given template name
func (builder *SiteBuilder) templatePath(tplName string) string {
	return path.Join(builder.themeTemplatesDir(), fmt.Sprintf("%s.html", tplName))
}

// Returns partials directory path
func (builder *SiteBuilder) partialsPath() string {
	return path.Join(builder.themeTemplatesDir(), PARTIALS_DIR)
}

func (builder *SiteBuilder) setupLayout() *template.Template {
	errStep := "template init"

	// parse layout
	result, err := template.ParseFiles(builder.templatePath("layout"))
	if err != nil {
		builder.addError(errStep, err)
		return nil
	}

	// setup FuncMap
	result.Funcs(builder.FuncMap())

	// load partials
	partialDir := builder.partialsPath()

	files, err := ioutil.ReadDir(partialDir)
	if err != nil && err != os.ErrExist {
		builder.addError(errStep, err)
	} else {
		for _, file := range files {
			fileName := file.Name()

			if !file.IsDir() && strings.HasSuffix(fileName, ".html") {
				filePath := path.Join(partialDir, fileName)

				// read partial
				binData, err := ioutil.ReadFile(filePath)
				if err != nil {
					builder.addError(errStep, err)
				} else {
					// eg: partials/navbar
					tplName := fmt.Sprintf("%s/%s", PARTIALS_DIR, utils.FileBase(fileName))

					// add partial to layout
					_, err := result.New(tplName).Parse(string(binData))
					if err != nil {
						builder.addError(errStep, err)
					}
				}
			}
		}
	}

	return result
}

// Get master layout template
func (builder *SiteBuilder) layout() *template.Template {
	if builder.masterLayout == nil {
		builder.masterLayout = builder.setupLayout()
	}

	return builder.masterLayout
}

// Dump templates
func (builder *SiteBuilder) DumpLayout() {
	log.Printf("Layout templates:")
	for _, tpl := range builder.masterLayout.Templates() {
		log.Printf("  -> %s", tpl.Name())
	}
}

// Prune directories for given absolute dir path
func (builder *SiteBuilder) ensureDir(dirPath string) error {
	// log.Printf("[DBG] Creating dir: %s", dirPath)

	err := os.MkdirAll(dirPath, 0777)
	if err != nil && err != os.ErrExist {
		return err
	}

	return err
}

// Prune directories for given absolute file path
func (builder *SiteBuilder) ensureFileDir(osPath string) error {
	return builder.ensureDir(path.Dir(osPath))
}

// Computes local file path for given URL
func (builder *SiteBuilder) filePath(fullUrl string) string {
	return path.Join(builder.GenDir(), fullUrl)
}

// Copy file to given directory
func (builder *SiteBuilder) copyFile(fromFilePath string, toDir string) error {
	// open source file
	src, err := os.Open(fromFilePath)
	if err != nil {
		return err
	}
	defer src.Close()

	// open destination file
	dstFilePath := path.Join(toDir, path.Base(fromFilePath))

	dst, err := os.Create(dstFilePath)
	if err != nil {
		return err
	}
	defer dst.Close()

	// copy
	if _, err := io.Copy(dst, src); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package simplestreams

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/osarch"
)

var urlDefaultOS = map[string]string{
	"https://cloud-images.ubuntu.com": "ubuntu",
}

// DownloadableFile represents a file with its URL, hash and size
type DownloadableFile struct {
	Path   string
	Sha256 string
	Size   int64
}

// NewClient returns a simplestreams client for the provided stream URL
func NewClient(url string, httpClient http.Client, useragent string) *SimpleStreams {
	return &SimpleStreams{
		http:           &httpClient,
		url:            url,
		cachedProducts: map[string]*Products{},
		useragent:      useragent,
	}
}

// SimpleStreams represents a simplestream client
type SimpleStreams struct {
	http      *http.Client
	url       string
	useragent string

	cachedStream   *Stream
	cachedProducts map[string]*Products
	cachedImages   []api.Image
	cachedAliases  []extendedAlias

	cachePath   string
	cacheExpiry time.Duration
}

// SetCache configures the on-disk cache
func (s *SimpleStreams) SetCache(path string, expiry time.Duration) {
	s.cachePath = path
	s.cacheExpiry = expiry
}

func (s *SimpleStreams) readCache(path string) ([]byte, bool) {
	cacheName := filepath.Join(s.cachePath, path)

	if s.cachePath == "" {
		return nil, false
	}

	if !shared.PathExists(cacheName) {
		return nil, false
	}

	fi, err := os.Stat(cacheName)
	if err != nil {
		os.Remove(cacheName)
		return nil, false
	}

	body, err := ioutil.ReadFile(cacheName)
	if err != nil {
		os.Remove(cacheName)
		return nil, false
	}

	expired := time.Now().Sub(fi.ModTime()) > s.cacheExpiry

	return body, expired
}

func (s *SimpleStreams) cachedDownload(path string) ([]byte, error) {
	fields := strings.Split(path, "/")
	fileName := fields[len(fields)-1]

	// Attempt to get from the cache
	cachedBody, expired := s.readCache(fileName)
	if cachedBody != nil && !expired {
		return cachedBody, nil
	}

	// Download from the source
	uri := fmt.Sprintf("%s/%s", s.url, path)
	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return nil, err
	}

	if s.useragent != "" {
		req.Header.Set("User-Agent", s.useragent)
	}

	r, err := s.http.Do(req)
	if err != nil {
		// On local connectivity error, return from cache anyway
		if cachedBody != nil {
			return cachedBody, nil
		}

		return nil, err
	}
	defer r.Body.Close()

	if r.StatusCode != http.StatusOK {
		// On local connectivity error, return from cache anyway
		if cachedBody != nil {
			return cachedBody, nil
		}

		return nil, fmt.Errorf("Unable to fetch %s: %s", uri, r.Status)
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}

	// Attempt to store in cache
	if s.cachePath != "" {
		cacheName := filepath.Join(s.cachePath, fileName)
		os.Remove(cacheName)
		ioutil.WriteFile(cacheName, body, 0644)
	}

	return body, nil
}

func (s *SimpleStreams) parseStream() (*Stream, error) {
	if s.cachedStream != nil {
		return s.cachedStream, nil
	}

	body, err := s.cachedDownload("/streams/v1/index.json")
	if err != nil {
		return nil, err
	}

	// Parse the idnex
	stream := Stream{}
	err = json.Unmarshal(body, &stream)
	if err != nil {
		return nil, err
	}

	s.cachedStream = &stream

	return &stream, nil
}

func (s *SimpleStreams) parseProducts(path string) (*Products, error) {
	if s.cachedProducts[path] != nil {
		return s.cachedProducts[path], nil
	}

	body, err := s.cachedDownload(path)
	if err != nil {
		return nil, err
	}

	// Parse the idnex
	products := Products{}
	err = json.Unmarshal(body, &products)
	if err != nil {
		return nil, err
	}

	s.cachedProducts[path] = &products

	return &products, nil
}

type extendedAlias struct {
	Name         string
	Alias        *api.ImageAliasesEntry
	Type         string
	Architecture string
}

func (s *SimpleStreams) applyAliases(images []api.Image) ([]api.Image, []extendedAlias, error) {
	aliasesList := []extendedAlias{}

	// Sort the images so we tag the preferred ones
	sort.Sort(sortedImages(images))

	// Look for the default OS
	defaultOS := ""
	for k, v := range urlDefaultOS {
		if strings.HasPrefix(s.url, k) {
			defaultOS = v
			break
		}
	}

	addAlias := func(imageType string, architecture string, name string, fingerprint string) *api.ImageAlias {
		if defaultOS != "" {
			name = strings.TrimPrefix(name, fmt.Sprintf("%s/", defaultOS))
		}

		for _, entry := range aliasesList {
			if entry.Name == name && entry.Type == imageType && entry.Architecture == architecture {
				return nil
			}
		}

		alias := api.ImageAliasesEntry{}
		alias.Name = name
		alias.Target = fingerprint
		alias.Type = imageType

		entry := extendedAlias{
			Name:         name,
			Type:         imageType,
			Alias:        &alias,
			Architecture: architecture,
		}

		aliasesList = append(aliasesList, entry)

		return &api.ImageAlias{Name: name}
	}

	architectureName, _ := osarch.ArchitectureGetLocal()

	newImages := []api.Image{}
	for _, image := range images {
		if image.Aliases != nil {
			// Build a new list of aliases from the provided ones
			aliases := image.Aliases
			image.Aliases = nil

			for _, entry := range aliases {
				// Short
				alias := addAlias(image.Type, image.Architecture, fmt.Sprintf("%s", entry.Name), image.Fingerprint)
				if alias != nil && architectureName == image.Architecture {
					image.Aliases = append(image.Aliases, *alias)
				}

				// Medium
				alias = addAlias(image.Type, image.Architecture, fmt.Sprintf("%s/%s", entry.Name, image.Properties["architecture"]), image.Fingerprint)
				if alias != nil {
					image.Aliases = append(image.Aliases, *alias)
				}
			}
		}

		newImages = append(newImages, image)
	}

	return newImages, aliasesList, nil
}

func (s *SimpleStreams) getImages() ([]api.Image, []extendedAlias, error) {
	if s.cachedImages != nil && s.cachedAliases != nil {
		return s.cachedImages, s.cachedAliases, nil
	}

	images := []api.Image{}

	// Load the stream data
	stream, err := s.parseStream()
	if err != nil {
		return nil, nil, err
	}

	// Iterate through the various indices
	for _, entry := range stream.Index {
		// We only care about images
		if entry.DataType != "image-downloads" {
			continue
		}

		// No point downloading an empty image list
		if len(entry.Products) == 0 {
			continue
		}

		products, err := s.parseProducts(entry.Path)
		if err != nil {
			return nil, nil, err
		}

		streamImages, _ := products.ToLXD()

		for _, image := range streamImages {
			images = append(images, image)
		}
	}

	// Setup the aliases
	images, aliases, err := s.applyAliases(images)
	if err != nil {
		return nil, nil, err
	}

	s.cachedImages = images
	s.cachedAliases = aliases

	return images, aliases, nil
}

// GetFiles returns a map of files for the provided image fingerprint
func (s *SimpleStreams) GetFiles(fingerprint string) (map[string]DownloadableFile, error) {
	// Load the main stream
	stream, err := s.parseStream()
	if err != nil {
		return nil, err
	}

	// Iterate through the various indices
	for _, entry := range stream.Index {
		// We only care about images
		if entry.DataType != "image-downloads" {
			continue
		}

		// No point downloading an empty image list
		if len(entry.Products) == 0 {
			continue
		}

		products, err := s.parseProducts(entry.Path)
		if err != nil {
			return nil, err
		}

		images, downloads := products.ToLXD()

		for _, image := range images {
			if strings.HasPrefix(image.Fingerprint, fingerprint) {
				files := map[string]DownloadableFile{}

				for _, path := range downloads[image.Fingerprint] {
					size, err := strconv.ParseInt(path[3], 10, 64)
					if err != nil {
						return nil, err
					}

					files[path[2]] = DownloadableFile{
						Path:   path[0],
						Sha256: path[1],
						Size:   size}
				}

				return files, nil
			}
		}
	}

	return nil, fmt.Errorf("Couldn't find the requested image")
}

// ListAliases returns a list of image aliases for the provided image fingerprint
func (s *SimpleStreams) ListAliases() ([]api.ImageAliasesEntry, error) {
	_, aliasesList, err := s.getImages()
	if err != nil {
		return nil, err
	}

	// Sort the list ahead of dedup
	sort.Sort(sortedAliases(aliasesList))

	aliases := []api.ImageAliasesEntry{}
	for _, entry := range aliasesList {
		dup := false
		for _, v := range aliases {
			if v.Name == entry.Name && v.Type == entry.Type {
				dup = true
			}
		}

		if dup {
			continue
		}

		aliases = append(aliases, *entry.Alias)
	}

	return aliases, nil
}

// ListImages returns a list of LXD images
func (s *SimpleStreams) ListImages() ([]api.Image, error) {
	images, _, err := s.getImages()
	return images, err
}

// GetAlias returns a LXD ImageAliasesEntry for the provided alias name
func (s *SimpleStreams) GetAlias(imageType string, name string) (*api.ImageAliasesEntry, error) {
	_, aliasesList, err := s.getImages()
	if err != nil {
		return nil, err
	}

	// Sort the list ahead of dedup
	sort.Sort(sortedAliases(aliasesList))

	var match *api.ImageAliasesEntry
	for _, entry := range aliasesList {
		if entry.Name != name {
			continue
		}

		if entry.Type != imageType && imageType != "" {
			continue
		}

		if match != nil {
			if match.Type != entry.Type {
				return nil, fmt.Errorf("More than one match for alias '%s'", name)
			}

			continue
		}

		match = entry.Alias
	}

	if match == nil {
		return nil, fmt.Errorf("Alias '%s' doesn't exist", name)
	}

	return match, nil
}

// GetAliasArchitectures returns a map of architecture / alias entries for an alias
func (s *SimpleStreams) GetAliasArchitectures(imageType string, name string) (map[string]*api.ImageAliasesEntry, error) {
	aliases := map[string]*api.ImageAliasesEntry{}

	_, aliasesList, err := s.getImages()
	if err != nil {
		return nil, err
	}

	for _, entry := range aliasesList {
		if entry.Name != name {
			continue
		}

		if entry.Type != imageType && imageType != "" {
			continue
		}

		if aliases[entry.Architecture] != nil {
			return nil, fmt.Errorf("More than one match for alias '%s'", name)
		}

		aliases[entry.Architecture] = entry.Alias
	}

	if len(aliases) == 0 {
		return nil, fmt.Errorf("Alias '%s' doesn't exist", name)
	}

	return aliases, nil
}

// GetImage returns a LXD image for the provided image fingerprint
func (s *SimpleStreams) GetImage(fingerprint string) (*api.Image, error) {
	images, _, err := s.getImages()
	if err != nil {
		return nil, err
	}

	matches := []api.Image{}

	for _, image := range images {
		if strings.HasPrefix(image.Fingerprint, fingerprint) {
			matches = append(matches, image)
		}
	}

	if len(matches) == 0 {
		return nil, fmt.Errorf("The requested image couldn't be found")
	} else if len(matches) > 1 {
		return nil, fmt.Errorf("More than one match for the provided partial fingerprint")
	}

	return &matches[0], nil
}
<commit_msg>shared/simplestreams: Fix stream's index download url<commit_after>package simplestreams

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/osarch"
)

var urlDefaultOS = map[string]string{
	"https://cloud-images.ubuntu.com": "ubuntu",
}

// DownloadableFile represents a file with its URL, hash and size
type DownloadableFile struct {
	Path   string
	Sha256 string
	Size   int64
}

// NewClient returns a simplestreams client for the provided stream URL
func NewClient(url string, httpClient http.Client, useragent string) *SimpleStreams {
	return &SimpleStreams{
		http:           &httpClient,
		url:            url,
		cachedProducts: map[string]*Products{},
		useragent:      useragent,
	}
}

// SimpleStreams represents a simplestream client
type SimpleStreams struct {
	http      *http.Client
	url       string
	useragent string

	cachedStream   *Stream
	cachedProducts map[string]*Products
	cachedImages   []api.Image
	cachedAliases  []extendedAlias

	cachePath   string
	cacheExpiry time.Duration
}

// SetCache configures the on-disk cache
func (s *SimpleStreams) SetCache(path string, expiry time.Duration) {
	s.cachePath = path
	s.cacheExpiry = expiry
}

func (s *SimpleStreams) readCache(path string) ([]byte, bool) {
	cacheName := filepath.Join(s.cachePath, path)

	if s.cachePath == "" {
		return nil, false
	}

	if !shared.PathExists(cacheName) {
		return nil, false
	}

	fi, err := os.Stat(cacheName)
	if err != nil {
		os.Remove(cacheName)
		return nil, false
	}

	body, err := ioutil.ReadFile(cacheName)
	if err != nil {
		os.Remove(cacheName)
		return nil, false
	}

	expired := time.Now().Sub(fi.ModTime()) > s.cacheExpiry

	return body, expired
}

func (s *SimpleStreams) cachedDownload(path string) ([]byte, error) {
	fields := strings.Split(path, "/")
	fileName := fields[len(fields)-1]

	// Attempt to get from the cache
	cachedBody, expired := s.readCache(fileName)
	if cachedBody != nil && !expired {
		return cachedBody, nil
	}

	// Download from the source
	uri := fmt.Sprintf("%s/%s", s.url, path)
	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return nil, err
	}

	if s.useragent != "" {
		req.Header.Set("User-Agent", s.useragent)
	}

	r, err := s.http.Do(req)
	if err != nil {
		// On local connectivity error, return from cache anyway
		if cachedBody != nil {
			return cachedBody, nil
		}

		return nil, err
	}
	defer r.Body.Close()

	if r.StatusCode != http.StatusOK {
		// On local connectivity error, return from cache anyway
		if cachedBody != nil {
			return cachedBody, nil
		}

		return nil, fmt.Errorf("Unable to fetch %s: %s", uri, r.Status)
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}

	// Attempt to store in cache
	if s.cachePath != "" {
		cacheName := filepath.Join(s.cachePath, fileName)
		os.Remove(cacheName)
		ioutil.WriteFile(cacheName, body, 0644)
	}

	return body, nil
}

func (s *SimpleStreams) parseStream() (*Stream, error) {
	if s.cachedStream != nil {
		return s.cachedStream, nil
	}

	body, err := s.cachedDownload("streams/v1/index.json")
	if err != nil {
		return nil, err
	}

	// Parse the idnex
	stream := Stream{}
	err = json.Unmarshal(body, &stream)
	if err != nil {
		return nil, err
	}

	s.cachedStream = &stream

	return &stream, nil
}

func (s *SimpleStreams) parseProducts(path string) (*Products, error) {
	if s.cachedProducts[path] != nil {
		return s.cachedProducts[path], nil
	}

	body, err := s.cachedDownload(path)
	if err != nil {
		return nil, err
	}

	// Parse the idnex
	products := Products{}
	err = json.Unmarshal(body, &products)
	if err != nil {
		return nil, err
	}

	s.cachedProducts[path] = &products

	return &products, nil
}

type extendedAlias struct {
	Name         string
	Alias        *api.ImageAliasesEntry
	Type         string
	Architecture string
}

func (s *SimpleStreams) applyAliases(images []api.Image) ([]api.Image, []extendedAlias, error) {
	aliasesList := []extendedAlias{}

	// Sort the images so we tag the preferred ones
	sort.Sort(sortedImages(images))

	// Look for the default OS
	defaultOS := ""
	for k, v := range urlDefaultOS {
		if strings.HasPrefix(s.url, k) {
			defaultOS = v
			break
		}
	}

	addAlias := func(imageType string, architecture string, name string, fingerprint string) *api.ImageAlias {
		if defaultOS != "" {
			name = strings.TrimPrefix(name, fmt.Sprintf("%s/", defaultOS))
		}

		for _, entry := range aliasesList {
			if entry.Name == name && entry.Type == imageType && entry.Architecture == architecture {
				return nil
			}
		}

		alias := api.ImageAliasesEntry{}
		alias.Name = name
		alias.Target = fingerprint
		alias.Type = imageType

		entry := extendedAlias{
			Name:         name,
			Type:         imageType,
			Alias:        &alias,
			Architecture: architecture,
		}

		aliasesList = append(aliasesList, entry)

		return &api.ImageAlias{Name: name}
	}

	architectureName, _ := osarch.ArchitectureGetLocal()

	newImages := []api.Image{}
	for _, image := range images {
		if image.Aliases != nil {
			// Build a new list of aliases from the provided ones
			aliases := image.Aliases
			image.Aliases = nil

			for _, entry := range aliases {
				// Short
				alias := addAlias(image.Type, image.Architecture, fmt.Sprintf("%s", entry.Name), image.Fingerprint)
				if alias != nil && architectureName == image.Architecture {
					image.Aliases = append(image.Aliases, *alias)
				}

				// Medium
				alias = addAlias(image.Type, image.Architecture, fmt.Sprintf("%s/%s", entry.Name, image.Properties["architecture"]), image.Fingerprint)
				if alias != nil {
					image.Aliases = append(image.Aliases, *alias)
				}
			}
		}

		newImages = append(newImages, image)
	}

	return newImages, aliasesList, nil
}

func (s *SimpleStreams) getImages() ([]api.Image, []extendedAlias, error) {
	if s.cachedImages != nil && s.cachedAliases != nil {
		return s.cachedImages, s.cachedAliases, nil
	}

	images := []api.Image{}

	// Load the stream data
	stream, err := s.parseStream()
	if err != nil {
		return nil, nil, err
	}

	// Iterate through the various indices
	for _, entry := range stream.Index {
		// We only care about images
		if entry.DataType != "image-downloads" {
			continue
		}

		// No point downloading an empty image list
		if len(entry.Products) == 0 {
			continue
		}

		products, err := s.parseProducts(entry.Path)
		if err != nil {
			return nil, nil, err
		}

		streamImages, _ := products.ToLXD()

		for _, image := range streamImages {
			images = append(images, image)
		}
	}

	// Setup the aliases
	images, aliases, err := s.applyAliases(images)
	if err != nil {
		return nil, nil, err
	}

	s.cachedImages = images
	s.cachedAliases = aliases

	return images, aliases, nil
}

// GetFiles returns a map of files for the provided image fingerprint
func (s *SimpleStreams) GetFiles(fingerprint string) (map[string]DownloadableFile, error) {
	// Load the main stream
	stream, err := s.parseStream()
	if err != nil {
		return nil, err
	}

	// Iterate through the various indices
	for _, entry := range stream.Index {
		// We only care about images
		if entry.DataType != "image-downloads" {
			continue
		}

		// No point downloading an empty image list
		if len(entry.Products) == 0 {
			continue
		}

		products, err := s.parseProducts(entry.Path)
		if err != nil {
			return nil, err
		}

		images, downloads := products.ToLXD()

		for _, image := range images {
			if strings.HasPrefix(image.Fingerprint, fingerprint) {
				files := map[string]DownloadableFile{}

				for _, path := range downloads[image.Fingerprint] {
					size, err := strconv.ParseInt(path[3], 10, 64)
					if err != nil {
						return nil, err
					}

					files[path[2]] = DownloadableFile{
						Path:   path[0],
						Sha256: path[1],
						Size:   size}
				}

				return files, nil
			}
		}
	}

	return nil, fmt.Errorf("Couldn't find the requested image")
}

// ListAliases returns a list of image aliases for the provided image fingerprint
func (s *SimpleStreams) ListAliases() ([]api.ImageAliasesEntry, error) {
	_, aliasesList, err := s.getImages()
	if err != nil {
		return nil, err
	}

	// Sort the list ahead of dedup
	sort.Sort(sortedAliases(aliasesList))

	aliases := []api.ImageAliasesEntry{}
	for _, entry := range aliasesList {
		dup := false
		for _, v := range aliases {
			if v.Name == entry.Name && v.Type == entry.Type {
				dup = true
			}
		}

		if dup {
			continue
		}

		aliases = append(aliases, *entry.Alias)
	}

	return aliases, nil
}

// ListImages returns a list of LXD images
func (s *SimpleStreams) ListImages() ([]api.Image, error) {
	images, _, err := s.getImages()
	return images, err
}

// GetAlias returns a LXD ImageAliasesEntry for the provided alias name
func (s *SimpleStreams) GetAlias(imageType string, name string) (*api.ImageAliasesEntry, error) {
	_, aliasesList, err := s.getImages()
	if err != nil {
		return nil, err
	}

	// Sort the list ahead of dedup
	sort.Sort(sortedAliases(aliasesList))

	var match *api.ImageAliasesEntry
	for _, entry := range aliasesList {
		if entry.Name != name {
			continue
		}

		if entry.Type != imageType && imageType != "" {
			continue
		}

		if match != nil {
			if match.Type != entry.Type {
				return nil, fmt.Errorf("More than one match for alias '%s'", name)
			}

			continue
		}

		match = entry.Alias
	}

	if match == nil {
		return nil, fmt.Errorf("Alias '%s' doesn't exist", name)
	}

	return match, nil
}

// GetAliasArchitectures returns a map of architecture / alias entries for an alias
func (s *SimpleStreams) GetAliasArchitectures(imageType string, name string) (map[string]*api.ImageAliasesEntry, error) {
	aliases := map[string]*api.ImageAliasesEntry{}

	_, aliasesList, err := s.getImages()
	if err != nil {
		return nil, err
	}

	for _, entry := range aliasesList {
		if entry.Name != name {
			continue
		}

		if entry.Type != imageType && imageType != "" {
			continue
		}

		if aliases[entry.Architecture] != nil {
			return nil, fmt.Errorf("More than one match for alias '%s'", name)
		}

		aliases[entry.Architecture] = entry.Alias
	}

	if len(aliases) == 0 {
		return nil, fmt.Errorf("Alias '%s' doesn't exist", name)
	}

	return aliases, nil
}

// GetImage returns a LXD image for the provided image fingerprint
func (s *SimpleStreams) GetImage(fingerprint string) (*api.Image, error) {
	images, _, err := s.getImages()
	if err != nil {
		return nil, err
	}

	matches := []api.Image{}

	for _, image := range images {
		if strings.HasPrefix(image.Fingerprint, fingerprint) {
			matches = append(matches, image)
		}
	}

	if len(matches) == 0 {
		return nil, fmt.Errorf("The requested image couldn't be found")
	} else if len(matches) > 1 {
		return nil, fmt.Errorf("More than one match for the provided partial fingerprint")
	}

	return &matches[0], nil
}
<|endoftext|>


<commit_before>package node

import (
	"fmt"
	"github.com/MG-RAST/Shock/shock-server/conf"
	e "github.com/MG-RAST/Shock/shock-server/errors"
	"github.com/MG-RAST/Shock/shock-server/logger"
	"github.com/MG-RAST/Shock/shock-server/node"
	"github.com/MG-RAST/Shock/shock-server/request"
	"github.com/MG-RAST/Shock/shock-server/responder"
	"github.com/MG-RAST/Shock/shock-server/util"
	"github.com/MG-RAST/golib/mgo/bson"
	"github.com/MG-RAST/golib/stretchr/goweb/context"
	"net/http"
	"strconv"
)

type M map[string]interface{}
type ListOfMaps []M

// GET: /node
// To do:
// - Iterate node queries
func (cr *NodeController) ReadMany(ctx context.Context) error {
	u, err := request.Authenticate(ctx.HttpRequest())
	if err != nil && err.Error() != e.NoAuth {
		return request.AuthError(err, ctx)
	}

	// Gather query params
	query := ctx.HttpRequest().URL.Query()

	// Setup query and nodes objects
	q := bson.M{}
	nodes := node.Nodes{}

	if u != nil {
		// Admin sees all
		if !u.Admin {
			q["$or"] = []bson.M{bson.M{"acl.read": []string{}}, bson.M{"acl.read": u.Uuid}, bson.M{"acl.owner": u.Uuid}}
		}
	} else {
		if conf.Bool(conf.Conf["anon-read"]) {
			// select on only nodes with no read rights set
			q["acl.read"] = []string{}
		} else {
			return responder.RespondWithError(ctx, http.StatusUnauthorized, e.NoAuth)
		}
	}

	// Gather params to make db query. Do not include the
	// following list.
	paramlist := map[string]int{"limit": 1, "offset": 1, "query": 1, "querynode": 1}
	if _, ok := query["query"]; ok {
		for key := range query {
			if _, found := paramlist[key]; !found {
				keyStr := fmt.Sprintf("attributes.%s", key)
				value := query.Get(key)
				if value != "" {
					if numValue, err := strconv.Atoi(value); err == nil {
						q["$or"] = ListOfMaps{{keyStr: value}, {keyStr: numValue}}
					} else {
						q[keyStr] = value
					}
				} else {
					existsMap := map[string]bool{
						"$exists": true,
					}
					q[keyStr] = existsMap
				}
			}
		}
	} else if _, ok := query["querynode"]; ok {
		for key := range query {
			if _, found := paramlist[key]; !found {
				value := query.Get(key)
				if value != "" {
					if numValue, err := strconv.Atoi(value); err == nil {
						q["$or"] = ListOfMaps{{key: value}, {key: numValue}}
					} else {
						q[key] = value
					}
				} else {
					existsMap := map[string]bool{
						"$exists": true,
					}
					q[key] = existsMap
				}
			}
		}
	}

	// defaults
	limit := 25
	offset := 0
	if _, ok := query["limit"]; ok {
		limit = util.ToInt(query.Get("limit"))
	}
	if _, ok := query["offset"]; ok {
		offset = util.ToInt(query.Get("offset"))
	}

	// Get nodes from db
	count, err := nodes.GetPaginated(q, limit, offset)
	if err != nil {
		err_msg := "err " + err.Error()
		logger.Error(err_msg)
		return responder.RespondWithError(ctx, http.StatusBadRequest, err_msg)
	}
	return responder.RespondWithPaginatedData(ctx, nodes, limit, offset, count)
}
<commit_msg>Added ability to query for field equal to null<commit_after>package node

import (
	"fmt"
	"github.com/MG-RAST/Shock/shock-server/conf"
	e "github.com/MG-RAST/Shock/shock-server/errors"
	"github.com/MG-RAST/Shock/shock-server/logger"
	"github.com/MG-RAST/Shock/shock-server/node"
	"github.com/MG-RAST/Shock/shock-server/request"
	"github.com/MG-RAST/Shock/shock-server/responder"
	"github.com/MG-RAST/Shock/shock-server/util"
	"github.com/MG-RAST/golib/mgo/bson"
	"github.com/MG-RAST/golib/stretchr/goweb/context"
	"net/http"
	"strconv"
)

type M map[string]interface{}
type ListOfMaps []M

// GET: /node
// To do:
// - Iterate node queries
func (cr *NodeController) ReadMany(ctx context.Context) error {
	u, err := request.Authenticate(ctx.HttpRequest())
	if err != nil && err.Error() != e.NoAuth {
		return request.AuthError(err, ctx)
	}

	// Gather query params
	query := ctx.HttpRequest().URL.Query()

	// Setup query and nodes objects
	q := bson.M{}
	nodes := node.Nodes{}

	if u != nil {
		// Admin sees all
		if !u.Admin {
			q["$or"] = []bson.M{bson.M{"acl.read": []string{}}, bson.M{"acl.read": u.Uuid}, bson.M{"acl.owner": u.Uuid}}
		}
	} else {
		if conf.Bool(conf.Conf["anon-read"]) {
			// select on only nodes with no read rights set
			q["acl.read"] = []string{}
		} else {
			return responder.RespondWithError(ctx, http.StatusUnauthorized, e.NoAuth)
		}
	}

	// Gather params to make db query. Do not include the
	// following list.
	paramlist := map[string]int{"limit": 1, "offset": 1, "query": 1, "querynode": 1}
	if _, ok := query["query"]; ok {
		for key := range query {
			if _, found := paramlist[key]; !found {
				keyStr := fmt.Sprintf("attributes.%s", key)
				value := query.Get(key)
				if value != "" {
					if numValue, err := strconv.Atoi(value); err == nil {
						q["$or"] = ListOfMaps{{keyStr: value}, {keyStr: numValue}}
					} else if value == "null" {
						q["$or"] = ListOfMaps{{keyStr: value}, {keyStr: nil}}
					} else {
						q[keyStr] = value
					}
				} else {
					existsMap := map[string]bool{
						"$exists": true,
					}
					q[keyStr] = existsMap
				}
			}
		}
	} else if _, ok := query["querynode"]; ok {
		for key := range query {
			if _, found := paramlist[key]; !found {
				value := query.Get(key)
				if value != "" {
					if numValue, err := strconv.Atoi(value); err == nil {
						q["$or"] = ListOfMaps{{key: value}, {key: numValue}}
					} else if value == "null" {
						q["$or"] = ListOfMaps{{key: value}, {key: nil}}
					} else {
						q[key] = value
					}
				} else {
					existsMap := map[string]bool{
						"$exists": true,
					}
					q[key] = existsMap
				}
			}
		}
	}

	// defaults
	limit := 25
	offset := 0
	if _, ok := query["limit"]; ok {
		limit = util.ToInt(query.Get("limit"))
	}
	if _, ok := query["offset"]; ok {
		offset = util.ToInt(query.Get("offset"))
	}

	// Get nodes from db
	count, err := nodes.GetPaginated(q, limit, offset)
	if err != nil {
		err_msg := "err " + err.Error()
		logger.Error(err_msg)
		return responder.RespondWithError(ctx, http.StatusBadRequest, err_msg)
	}
	return responder.RespondWithPaginatedData(ctx, nodes, limit, offset, count)
}
<|endoftext|>


<commit_before>package scrape

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/segmentio/ksuid"
	"github.com/slotix/dataflowkit/errs"
	"github.com/slotix/dataflowkit/extract"
	"github.com/slotix/dataflowkit/fetch"
	"github.com/slotix/dataflowkit/paginate"
	"github.com/slotix/dataflowkit/splash"
	"github.com/spf13/viper"
)

var logger *log.Logger

func init() {
	logger = log.New(os.Stdout, "scrape: ", log.Lshortfile)
}

var (
	ErrNoParts     = errors.New("no pieces in payload")
	ErrNoSelectors = errors.New("no selectors in payload")
)

// First returns the first set of results - i.e. the results from the first
// block on the first page.
// This function can return nil if there were no blocks found on the first page
// of the scrape.
func (r *Results) First() map[string]interface{} {
	if len(r.Output[0]) == 0 {
		return nil
	}

	return r.Output[0][0]
}

// AllBlocks returns a single list of results from every block on all pages.
// This function will always return a list, even if no blocks were found.
func (r *Results) AllBlocks() []map[string]interface{} {
	ret := []map[string]interface{}{}

	for _, page := range r.Output {
		for _, block := range page {
			ret = append(ret, block)
		}
	}

	return ret
}

//KSUID stores the timestamp portion in ID. So we can retrieve it from Task object as a Time object
func (t Task) StartTime() (*time.Time, error) {
	id, err := ksuid.Parse(t.ID)
	if err != nil {
		return nil, err
	}
	idTime := id.Time()
	return &idTime, nil
}

//selectors returns selectors from payload
func (p Payload) selectors() ([]string, error) {
	selectors := []string{}
	for _, f := range p.Fields {
		selectors = append(selectors, f.Selector)
	}
	if len(selectors) == 0 {
		return nil, errNoSelectors
	}
	return selectors, nil
}

func NewTask(p Payload) (task *Task, err error) {
	scraper := &Scraper{}
	//if p.Request != nil {
		scraper, err = NewScraper(p)
		if err != nil {
			return nil, err
		}
//	}
	//https://blog.kowalczyk.info/article/JyRZ/generating-good-random-and-unique-ids-in-go.html
	id := ksuid.New()

	task = &Task{
		ID:      id.String(),
		Scraper: scraper,
	}
	return task, nil
}

//fields2parts converts payload []field to []scrape.Part
func (p Payload) fields2parts() ([]Part, error) {
	parts := []Part{}
	//Payload fields
	for _, f := range p.Fields {
		params := make(map[string]interface{})
		if f.Extractor.Params != nil {
			params = f.Extractor.Params.(map[string]interface{})
		}
		var err error
		switch eType := f.Extractor.Type; eType {

		//For Link type Two pieces as pair Text and Attr{Attr:"href"} extractors are added.
		case "link":
			l := &extract.Link{Href: extract.Attr{Attr: "href"}}
			if params != nil {
				err := FillStruct(params, l)
				if err != nil {
					logger.Println(err)
				}
			}
			//******* details
			task := &Task{}
			if f.Details != nil {
				detailsPayload := p
				detailsPayload.Name = f.Name + "Details"
				detailsPayload.Fields = f.Details.Fields
				//Request refers to  srarting URL here. Requests will be changed in Scrape function to Details pages afterwards
				task, err = NewTask(detailsPayload)
				if err != nil {
					return nil, err
				}
			}
			parts = append(parts, Part{
				Name:      f.Name + "_text",
				Selector:  f.Selector,
				Extractor: l.Text,
			}, Part{
				Name:      f.Name + "_link",
				Selector:  f.Selector,
				Extractor: l.Href,
				Details:   task,
			})

		//For image type by default Two pieces with different Attr="src" and Attr="alt" extractors will be added for field selector.
		case "image":
			i := &extract.Image{Src: extract.Attr{Attr: "src"},
				Alt: extract.Attr{Attr: "alt"}}
			if params != nil {
				err := FillStruct(params, i)
				if err != nil {
					logger.Println(err)
				}
			}
			parts = append(parts, Part{
				Name:      f.Name + "_src",
				Selector:  f.Selector,
				Extractor: i.Src,
			}, Part{
				Name:      f.Name + "_alt",
				Selector:  f.Selector,
				Extractor: i.Alt,
			})

		default:
			var e extract.Extractor
			switch eType {
			case "const":
				//	c := &extract.Const{Val: params["value"]}
				//	e = c
				e = &extract.Const{}
			case "count":
				e = &extract.Count{}
			case "text":
				e = &extract.Text{}
			case "html":
				e = &extract.Html{}
			case "outerHtml":
				e = &extract.OuterHtml{}
			case "attr":
				e = &extract.Attr{}
			case "regex":
				r := &extract.Regex{}
				regExp := params["regexp"]
				r.Regex = regexp.MustCompile(regExp.(string))
				e = r
			}

			if params != nil {
				err := FillStruct(params, e)
				if err != nil {
					logger.Println(err)
				}
			}
			parts = append(parts, Part{
				Name:      f.Name,
				Selector:  f.Selector,
				Extractor: e,
			})
		}
	}
	// Validate payload fields
	if len(parts) == 0 {
		return nil, ErrNoParts
	}
	seenNames := map[string]struct{}{}
	for i, part := range parts {
		if len(part.Name) == 0 {
			return nil, fmt.Errorf("no name provided for part %d", i)
		}
		if _, seen := seenNames[part.Name]; seen {
			return nil, fmt.Errorf("part %s has a duplicate name", i)
		}
		seenNames[part.Name] = struct{}{}

		if len(part.Selector) == 0 {
			return nil, fmt.Errorf("no selector provided for part %d", i)
		}
	}
	return parts, nil
}

// Create a new scraper with the provided configuration.
func NewScraper(p Payload) (*Scraper, error) {
	parts, err := p.fields2parts()
	if err != nil {
		return nil, err
	}
	var paginator paginate.Paginator
	if p.Paginator == nil {
		paginator = &dummyPaginator{}
	} else {
		paginator = paginate.BySelector(p.Paginator.Selector, p.Paginator.Attribute)
	}

	selectors, err := p.selectors()
	if err != nil {
		return nil, err
	}
	//TODO: need to test the case when there are no selectors found in payload.
	var dividePageFunc DividePageFunc
	if len(selectors) == 0 {
		dividePageFunc = DividePageBySelector("body")
	} else {
		dividePageFunc = DividePageByIntersection(selectors)
	}

	scraper := &Scraper{
		Request: p.Request.(splash.Request),
		DividePage: dividePageFunc,
		Parts:      parts,
		Paginator:  paginator,
		Opts: ScrapeOptions{
			MaxPages:            p.Paginator.MaxPages,
			Format:              p.Format,
			PaginateResults:     *p.PaginateResults,
			FetchDelay:          p.FetchDelay,
			RandomizeFetchDelay: *p.RandomizeFetchDelay,
			RetryTimes:          p.RetryTimes,
		},
	}

	// All set!
	return scraper, nil
}

func  Scrape(task *Task) error {
	req := task.Scraper.Request
	url := req.GetURL()

	var numPages int
	opts := task.Scraper.Opts
	task.Results.Visited = make(map[string]error)

	//get Robotstxt Data
	 robots, err := fetch.RobotstxtData(url)
	 if err != nil {
		task.Results.Visited[url] = err
		logger.Println(err)
		//return err
	 }
	 task.Scraper.Robots = robots

	for {
		//check if scraping of current url is not forbidden
		if !fetch.AllowedByRobots(url, task.Scraper.Robots) {
			task.Results.Visited[url] = &errs.ForbiddenByRobots{url}
		}
		// Repeat until we don't have any more URLs, or until we hit our page limit.
		if len(url) == 0 ||
			(opts.MaxPages > 0 && numPages >= opts.MaxPages) {
			break
		}
		//call remote fetcher to download web page
		sResponse, err := responseFromFetchService(req)
		//sResponse, err := req.(splash.Request).GetResponse()
		if err != nil {
			return err
		}
		content, err := sResponse.GetContent()
		if err != nil {
			return err
		}
		// Create a goquery document.
		doc, err := goquery.NewDocumentFromReader(content)
		if err != nil {
			return err
		}

		task.Results.Visited[url] = nil
		results := []map[string]interface{}{}

		// Divide this page into blocks
		for _, block := range task.Scraper.DividePage(doc.Selection) {
			blockResults := map[string]interface{}{}

			// Process each part of this block
			for _, part := range task.Scraper.Parts {
				sel := block
				if part.Selector != "." {
					sel = sel.Find(part.Selector)
				}

				partResults, err := part.Extractor.Extract(sel)
				if err != nil {
					return err
				}

				// A nil response from an extractor means that we don't even include it in
				// the results.
				if partResults == nil {
					continue
				}

				blockResults[part.Name] = partResults

				//********* details
				if part.Details != nil {
					//fmt.Printf("%T\n", partResults)
					part.Details.Scraper.Request = splash.Request{URL: partResults.(string)}
					err = Scrape(part.Details)
					if err != nil {
						return err
					}
					//logger.Println(part.Details.Visited)
					//logger.Println(part.Details.Results)
					//part.Details.Results = append(part.Details.Results, )
				//	logger.Println(blockResults[part.Name], part.Details)
				}
			}
			if len(blockResults) > 0 {
				// Append the results from this block.
				results = append(results, blockResults)
			}
		}

		// Append the results from this page.
		task.Results.Output = append(task.Results.Output, results)

		numPages++

		// Get the next page.
		url, err = task.Scraper.Paginator.NextPage(url, doc.Selection)
		if err != nil {
			return err
		}

		//every time when getting a response the next request will be filled with updated cookie information
		err = sResponse.SetCookieToNextRequest(&req)
		if err != nil {
			logger.Println(err)
		}
		req.URL = url

		if opts.RandomizeFetchDelay {
			//Sleep for time equal to FetchDelay * random value between 500 and 1500 msec
			rand := Random(500, 1500)
			delay := opts.FetchDelay * time.Duration(rand) / 1000
			logger.Println(delay)
			time.Sleep(delay)
		} else {
			time.Sleep(opts.FetchDelay)
		}

	}
	logger.Println(task.Visited)
	// All good!
	return nil

}

//responseFromFetchService sends request to fetch service and returns *splash.Response
func responseFromFetchService(req splash.Request) (*splash.Response, error) {

	//fetch content
	b, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}
	reader := bytes.NewReader(b)
	addr := "http://" + viper.GetString("DFK_FETCH") + "/response/splash"
	request, err := http.NewRequest("POST", addr, reader)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	r, err := client.Do(request)
	if r != nil {
		defer r.Body.Close()
	}
	if err != nil {
		panic(err)
	}
	resp, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}
	var sResponse *splash.Response
	if err := json.Unmarshal(resp, &sResponse); err != nil {
		logger.Println("Json Unmarshall error", err)
	}
	//	content, err := sResponse.GetContent()
	//	if err != nil {
	//		return nil, err
	//	}
	//return content, nil
	return sResponse, nil
}

//PartNames returns Part Names which are used as a header of output CSV
func (s Scraper) PartNames() []string {
	names := []string{}
	for _, part := range s.Parts {
		names = append(names, part.Name)
	}
	return names
}
<commit_msg>fixed paginator issue<commit_after>package scrape

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/segmentio/ksuid"
	"github.com/slotix/dataflowkit/errs"
	"github.com/slotix/dataflowkit/extract"
	"github.com/slotix/dataflowkit/fetch"
	"github.com/slotix/dataflowkit/paginate"
	"github.com/slotix/dataflowkit/splash"
	"github.com/spf13/viper"
)

var logger *log.Logger

func init() {
	logger = log.New(os.Stdout, "scrape: ", log.Lshortfile)
}

var (
	ErrNoParts     = errors.New("no pieces in payload")
	ErrNoSelectors = errors.New("no selectors in payload")
)

// First returns the first set of results - i.e. the results from the first
// block on the first page.
// This function can return nil if there were no blocks found on the first page
// of the scrape.
func (r *Results) First() map[string]interface{} {
	if len(r.Output[0]) == 0 {
		return nil
	}

	return r.Output[0][0]
}

// AllBlocks returns a single list of results from every block on all pages.
// This function will always return a list, even if no blocks were found.
func (r *Results) AllBlocks() []map[string]interface{} {
	ret := []map[string]interface{}{}

	for _, page := range r.Output {
		for _, block := range page {
			ret = append(ret, block)
		}
	}

	return ret
}

//KSUID stores the timestamp portion in ID. So we can retrieve it from Task object as a Time object
func (t Task) StartTime() (*time.Time, error) {
	id, err := ksuid.Parse(t.ID)
	if err != nil {
		return nil, err
	}
	idTime := id.Time()
	return &idTime, nil
}

//selectors returns selectors from payload
func (p Payload) selectors() ([]string, error) {
	selectors := []string{}
	for _, f := range p.Fields {
		selectors = append(selectors, f.Selector)
	}
	if len(selectors) == 0 {
		return nil, errNoSelectors
	}
	return selectors, nil
}

func NewTask(p Payload) (task *Task, err error) {
	scraper := &Scraper{}
	//if p.Request != nil {
		scraper, err = NewScraper(p)
		if err != nil {
			return nil, err
		}
//	}
	//https://blog.kowalczyk.info/article/JyRZ/generating-good-random-and-unique-ids-in-go.html
	id := ksuid.New()

	task = &Task{
		ID:      id.String(),
		Scraper: scraper,
	}
	return task, nil
}

//fields2parts converts payload []field to []scrape.Part
func (p Payload) fields2parts() ([]Part, error) {
	parts := []Part{}
	//Payload fields
	for _, f := range p.Fields {
		params := make(map[string]interface{})
		if f.Extractor.Params != nil {
			params = f.Extractor.Params.(map[string]interface{})
		}
		var err error
		switch eType := f.Extractor.Type; eType {

		//For Link type Two pieces as pair Text and Attr{Attr:"href"} extractors are added.
		case "link":
			l := &extract.Link{Href: extract.Attr{Attr: "href"}}
			if params != nil {
				err := FillStruct(params, l)
				if err != nil {
					logger.Println(err)
				}
			}
			parts = append(parts, Part{
				Name:      f.Name + "_text",
				Selector:  f.Selector,
				Extractor: l.Text,
			})
			//******* details

			if f.Details != nil {
				task := &Task{}
				detailsPayload := p
				detailsPayload.Name = f.Name + "Details"
				detailsPayload.Fields = f.Details.Fields
				//Request refers to  srarting URL here. Requests will be changed in Scrape function to Details pages afterwards
				task, err = NewTask(detailsPayload)
				if err != nil {
					return nil, err
				}
				parts = append(parts,  Part{
					Name:      f.Name + "_link",
					Selector:  f.Selector,
					Extractor: l.Href,
					Details:   task,
				})
			}
			parts = append(parts,  Part{
				Name:      f.Name + "_link",
				Selector:  f.Selector,
				Extractor: l.Href,
				Details:   nil,
			})

		//For image type by default Two pieces with different Attr="src" and Attr="alt" extractors will be added for field selector.
		case "image":
			i := &extract.Image{Src: extract.Attr{Attr: "src"},
				Alt: extract.Attr{Attr: "alt"}}
			if params != nil {
				err := FillStruct(params, i)
				if err != nil {
					logger.Println(err)
				}
			}
			parts = append(parts, Part{
				Name:      f.Name + "_src",
				Selector:  f.Selector,
				Extractor: i.Src,
			}, Part{
				Name:      f.Name + "_alt",
				Selector:  f.Selector,
				Extractor: i.Alt,
			})

		default:
			var e extract.Extractor
			switch eType {
			case "const":
				//	c := &extract.Const{Val: params["value"]}
				//	e = c
				e = &extract.Const{}
			case "count":
				e = &extract.Count{}
			case "text":
				e = &extract.Text{}
			case "html":
				e = &extract.Html{}
			case "outerHtml":
				e = &extract.OuterHtml{}
			case "attr":
				e = &extract.Attr{}
			case "regex":
				r := &extract.Regex{}
				regExp := params["regexp"]
				r.Regex = regexp.MustCompile(regExp.(string))
				e = r
			}

			if params != nil {
				err := FillStruct(params, e)
				if err != nil {
					logger.Println(err)
				}
			}
			parts = append(parts, Part{
				Name:      f.Name,
				Selector:  f.Selector,
				Extractor: e,
			})
		}
	}
	// Validate payload fields
	if len(parts) == 0 {
		return nil, ErrNoParts
	}
	seenNames := map[string]struct{}{}
	for i, part := range parts {
		if len(part.Name) == 0 {
			return nil, fmt.Errorf("no name provided for part %d", i)
		}
		if _, seen := seenNames[part.Name]; seen {
			return nil, fmt.Errorf("part %s has a duplicate name", i)
		}
		seenNames[part.Name] = struct{}{}

		if len(part.Selector) == 0 {
			return nil, fmt.Errorf("no selector provided for part %d", i)
		}
	}
	return parts, nil
}

// Create a new scraper with the provided configuration.
func NewScraper(p Payload) (*Scraper, error) {
	parts, err := p.fields2parts()
	if err != nil {
		return nil, err
	}
	var paginator paginate.Paginator
	maxPages := 1
	if p.Paginator == nil {
		paginator = &dummyPaginator{}
	} else {
		paginator = paginate.BySelector(p.Paginator.Selector, p.Paginator.Attribute)
		maxPages = p.Paginator.MaxPages
	}

	selectors, err := p.selectors()
	if err != nil {
		return nil, err
	}
	//TODO: need to test the case when there are no selectors found in payload.
	var dividePageFunc DividePageFunc
	if len(selectors) == 0 {
		dividePageFunc = DividePageBySelector("body")
	} else {
		dividePageFunc = DividePageByIntersection(selectors)
	}

	scraper := &Scraper{
		Request: p.Request.(splash.Request),
		DividePage: dividePageFunc,
		Parts:      parts,
		Paginator:  paginator,
		Opts: ScrapeOptions{
			MaxPages:            maxPages,
			Format:              p.Format,
			PaginateResults:     *p.PaginateResults,
			FetchDelay:          p.FetchDelay,
			RandomizeFetchDelay: *p.RandomizeFetchDelay,
			RetryTimes:          p.RetryTimes,
		},
	}

	// All set!
	return scraper, nil
}

func  Scrape(task *Task) error {
	req := task.Scraper.Request
	url := req.GetURL()

	var numPages int
	opts := task.Scraper.Opts
	task.Results.Visited = make(map[string]error)

	//get Robotstxt Data
	 robots, err := fetch.RobotstxtData(url)
	 if err != nil {
		task.Results.Visited[url] = err
		logger.Println(err)
		//return err
	 }
	 task.Scraper.Robots = robots

	for {
		//check if scraping of current url is not forbidden
		if !fetch.AllowedByRobots(url, task.Scraper.Robots) {
			task.Results.Visited[url] = &errs.ForbiddenByRobots{url}
		}
		// Repeat until we don't have any more URLs, or until we hit our page limit.
		if len(url) == 0 ||
			(opts.MaxPages > 0 && numPages >= opts.MaxPages) {
			break
		}
		//call remote fetcher to download web page
		sResponse, err := responseFromFetchService(req)
		//sResponse, err := req.(splash.Request).GetResponse()
		if err != nil {
			return err
		}
		content, err := sResponse.GetContent()
		if err != nil {
			return err
		}
		// Create a goquery document.
		doc, err := goquery.NewDocumentFromReader(content)
		if err != nil {
			return err
		}

		task.Results.Visited[url] = nil
		results := []map[string]interface{}{}

		// Divide this page into blocks
		for _, block := range task.Scraper.DividePage(doc.Selection) {
			blockResults := map[string]interface{}{}

			// Process each part of this block
			for _, part := range task.Scraper.Parts {
				sel := block
				if part.Selector != "." {
					sel = sel.Find(part.Selector)
				}

				partResults, err := part.Extractor.Extract(sel)
				if err != nil {
					return err
				}

				// A nil response from an extractor means that we don't even include it in
				// the results.
				if partResults == nil {
					continue
				}

				blockResults[part.Name] = partResults

				//********* details
				if part.Details != nil {
					//fmt.Printf("%T\n", partResults)
					part.Details.Scraper.Request = splash.Request{URL: partResults.(string)}
					err = Scrape(part.Details)
					if err != nil {
						return err
					}
					//logger.Println(part.Details.Visited)
					//logger.Println(part.Details.Results)
					//part.Details.Results = append(part.Details.Results, )
				//	logger.Println(blockResults[part.Name], part.Details)
				}
			}
			if len(blockResults) > 0 {
				// Append the results from this block.
				results = append(results, blockResults)
			}
		}

		// Append the results from this page.
		task.Results.Output = append(task.Results.Output, results)

		numPages++

		// Get the next page.
		url, err = task.Scraper.Paginator.NextPage(url, doc.Selection)
		if err != nil {
			return err
		}

		//every time when getting a response the next request will be filled with updated cookie information
		err = sResponse.SetCookieToNextRequest(&req)
		if err != nil {
			logger.Println(err)
		}
		req.URL = url

		if opts.RandomizeFetchDelay {
			//Sleep for time equal to FetchDelay * random value between 500 and 1500 msec
			rand := Random(500, 1500)
			delay := opts.FetchDelay * time.Duration(rand) / 1000
			logger.Println(delay)
			time.Sleep(delay)
		} else {
			time.Sleep(opts.FetchDelay)
		}

	}
	logger.Println(task.Visited)
	// All good!
	return nil

}

//responseFromFetchService sends request to fetch service and returns *splash.Response
func responseFromFetchService(req splash.Request) (*splash.Response, error) {

	//fetch content
	b, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}
	reader := bytes.NewReader(b)
	addr := "http://" + viper.GetString("DFK_FETCH") + "/response/splash"
	request, err := http.NewRequest("POST", addr, reader)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	r, err := client.Do(request)
	if r != nil {
		defer r.Body.Close()
	}
	if err != nil {
		panic(err)
	}
	resp, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}
	var sResponse *splash.Response
	if err := json.Unmarshal(resp, &sResponse); err != nil {
		logger.Println("Json Unmarshall error", err)
	}
	//	content, err := sResponse.GetContent()
	//	if err != nil {
	//		return nil, err
	//	}
	//return content, nil
	return sResponse, nil
}

//PartNames returns Part Names which are used as a header of output CSV
func (s Scraper) PartNames() []string {
	names := []string{}
	for _, part := range s.Parts {
		names = append(names, part.Name)
	}
	return names
}
<|endoftext|>


<commit_before>package transport

import (
	"github.com/AutogrowSystems/faye-go/protocol"
	"github.com/AutogrowSystems/faye-go/utils"
	"github.com/gorilla/websocket"
	"io"
)

const WebSocketConnectionPriority = 10

type Server interface {
	HandleRequest(interface{}, protocol.Connection)
	Logger() utils.Logger
}

type WebSocketConnection struct {
	ws         *websocket.Conn
	failedSend bool
}

func (wc *WebSocketConnection) Send(msgs []protocol.Message) error {
	err := wc.ws.WriteJSON(msgs)
	if err != nil {
		wc.failedSend = true
	}
	return err
}

func (wc *WebSocketConnection) IsConnected() bool {
	return wc.failedSend
}

func (wc *WebSocketConnection) Close() {
	wc.ws.Close()
}

func (wc WebSocketConnection) Priority() int {
	return WebSocketConnectionPriority
}

func (lp WebSocketConnection) IsSingleShot() bool {
	return false
}

func WebsocketServer(m Server) func(*websocket.Conn) {
	return func(ws *websocket.Conn) {
		var data interface{}
		wsConn := WebSocketConnection{ws, true}
		for {
			err := ws.ReadJSON(&data)
			if err != nil {
				if err == io.EOF {
					m.Logger().Debugf("EOF while reading from socket")
					return
				} else {
					m.Logger().Debugf("While reading from socket: %s", err)
					return
				}
			}

			if arr := data.([]interface{}); len(arr) == 0 {
				ws.WriteJSON([]string{})
			} else {
				m.HandleRequest(data, &wsConn)
			}
		}
	}
}
<commit_msg>changed logging level from debug to error for websocket errors<commit_after>package transport

import (
	"github.com/AutogrowSystems/faye-go/protocol"
	"github.com/AutogrowSystems/faye-go/utils"
	"github.com/gorilla/websocket"
	"io"
)

const WebSocketConnectionPriority = 10

type Server interface {
	HandleRequest(interface{}, protocol.Connection)
	Logger() utils.Logger
}

type WebSocketConnection struct {
	ws         *websocket.Conn
	failedSend bool
}

func (wc *WebSocketConnection) Send(msgs []protocol.Message) error {
	err := wc.ws.WriteJSON(msgs)
	if err != nil {
		wc.failedSend = true
	}
	return err
}

func (wc *WebSocketConnection) IsConnected() bool {
	return wc.failedSend
}

func (wc *WebSocketConnection) Close() {
	wc.ws.Close()
}

func (wc WebSocketConnection) Priority() int {
	return WebSocketConnectionPriority
}

func (lp WebSocketConnection) IsSingleShot() bool {
	return false
}

func WebsocketServer(m Server) func(*websocket.Conn) {
	return func(ws *websocket.Conn) {
		var data interface{}
		wsConn := WebSocketConnection{ws, true}
		for {
			err := ws.ReadJSON(&data)
			if err != nil {
				if err == io.EOF {
					m.Logger().Errorf("EOF while reading from socket")
					return
				} else {
					m.Logger().Errorf("While reading from socket: %s", err)
					return
				}
			}

			if arr := data.([]interface{}); len(arr) == 0 {
				ws.WriteJSON([]string{})
			} else {
				m.HandleRequest(data, &wsConn)
			}
		}
	}
}
<|endoftext|>


<commit_before>package datastore

import (
	"sort"

	"github.com/herb-go/herb/cache"
)

func unmarshalMapElement(s Store, creator func() interface{}, key string, data []byte, c cache.Cacheable) (err error) {
	v := creator()
	if len(data) != 0 {
		err = c.Unmarshal(data, v)
		if err != nil {
			return err
		}
	} else {
		v = nil
	}
	s.Store(key, v)
	return nil
}

//Load load data by given key list into data store.
//Param s target data store.
//Param c map cache
//Param loader func by which data load.
//Param creator map element creator.
//Param keys key list to load
//Return any error if raised.
func Load(s Store, c cache.Cacheable, loader func(...string) (map[string]interface{}, error), creator func() interface{}, keys ...string) error {
	var keysmap = make(map[string]bool, len(keys))
	var filteredKeys = make([]string, len(keys))
	var filteredKeysLength = 0
	var err error
	if len(keys) == 0 {
		return nil
	}
	for k := range keys {
		if keysmap[keys[k]] == true {
			continue
		}
		keysmap[keys[k]] = true
		if _, ok := s.Load(keys[k]); !ok {

			filteredKeys[filteredKeysLength] = keys[k]
			filteredKeysLength++
		}
	}
	filteredKeys = filteredKeys[:filteredKeysLength]
	if filteredKeysLength == 0 {
		return nil
	}
	var results map[string][]byte
	if c != nil {

		results, err = c.MGetBytesValue(filteredKeys...)
		if err != nil {
			return err
		}
	} else {
		results = map[string][]byte{}
	}
	var uncachedKeys = make([]string, len(filteredKeys))
	var uncachedKeysLength = 0
	for i := range filteredKeys {
		k := filteredKeys[i]
		if _, ok := results[k]; !ok {
			uncachedKeys[uncachedKeysLength] = k
			uncachedKeysLength++
		} else {
			err = unmarshalMapElement(s, creator, k, results[k], c)
			if err != nil {
				return err
			}
		}
	}
	uncachedKeys = uncachedKeys[:uncachedKeysLength]

	if uncachedKeysLength == 0 {
		return nil
	}
	var unreloadkeys = make([]string, len(uncachedKeys))
	var unreloadkeysLength = 0

	sort.Strings(uncachedKeys)
	if c != nil {
		for k := range uncachedKeys {
			key, err := c.FinalKey(uncachedKeys[k])
			if err != nil {
				return err
			}
			locker, ok := c.Locker(key)
			if ok {
				locker.RLock()
				defer locker.RUnlock()
				bs, err := c.GetBytesValue(uncachedKeys[k])
				if err == nil {
					err = unmarshalMapElement(s, creator, uncachedKeys[k], bs, c)
					if err != nil {
						return err
					}
					continue
				}
				if err != cache.ErrNotFound {
					return err
				}
			} else {
				locker.Lock()
				defer locker.Unlock()
			}
			unreloadkeys[unreloadkeysLength] = uncachedKeys[k]
			unreloadkeysLength++

		}

	} else {
		unreloadkeys = uncachedKeys
		unreloadkeysLength = uncachedKeysLength
	}
	unreloadkeys = unreloadkeys[:unreloadkeysLength]

	if unreloadkeysLength == 0 {
		return nil
	}
	loaded, err := loader(unreloadkeys...)
	if err != nil {
		return err
	}
	var data = make(map[string][]byte, len(loaded))
	for k := range loaded {
		v := loaded[k]
		s.Store(k, v)
		if c != nil {
			data[k], err = c.Marshal(v)
			if err != nil {
				return err
			}
		}
	}
	for k := range unreloadkeys {
		if _, ok := data[unreloadkeys[k]]; ok == false {
			data[unreloadkeys[k]] = []byte{}
		}
	}
	if c == nil {
		return nil
	}
	return c.MSetBytesValue(data, 0)
}
<commit_msg>update<commit_after>package datastore

import (
	"sort"

	"github.com/herb-go/herb/cache"
)

func unmarshalMapElement(s Store, creator func() interface{}, key string, data []byte, c cache.Cacheable) (err error) {
	v := creator()
	if len(data) != 0 {
		err = c.Unmarshal(data, v)
		if err != nil {
			return err
		}
	} else {
		v = nil
	}
	s.Store(key, v)
	return nil
}

//Load load data by given key list into data store.
//Param s target data store.
//Param c map cache
//Param loader func by which data load.
//Param creator map element creator.
//Param keys key list to load
//Return any error if raised.
func Load(s Store, c cache.Cacheable, loader func(...string) (map[string]interface{}, error), creator func() interface{}, keys ...string) error {
	var keysmap = make(map[string]bool, len(keys))
	var filteredKeys = make([]string, 0, len(keys))
	var err error
	if len(keys) == 0 {
		return nil
	}
	for k := range keys {
		if keysmap[keys[k]] == true {
			continue
		}
		keysmap[keys[k]] = true
		if _, ok := s.Load(keys[k]); !ok {

			filteredKeys = append(filteredKeys, keys[k])
		}
	}
	if len(filteredKeys) == 0 {
		return nil
	}
	var results map[string][]byte
	if c != nil {

		results, err = c.MGetBytesValue(filteredKeys...)
		if err != nil {
			return err
		}
	} else {
		results = map[string][]byte{}
	}
	var uncachedKeys = make([]string, 0, len(filteredKeys))
	for i := range filteredKeys {
		k := filteredKeys[i]
		if _, ok := results[k]; !ok {
			uncachedKeys = append(uncachedKeys, k)
		} else {
			err = unmarshalMapElement(s, creator, k, results[k], c)
			if err != nil {
				return err
			}
		}
	}
	if len(uncachedKeys) == 0 {
		return nil
	}
	var unreloadkeys = make([]string, 0, len(uncachedKeys))

	sort.Strings(uncachedKeys)
	if c != nil {
		for k := range uncachedKeys {
			key, err := c.FinalKey(uncachedKeys[k])
			if err != nil {
				return err
			}
			locker, ok := c.Locker(key)
			if ok {
				locker.RLock()
				defer locker.RUnlock()
				bs, err := c.GetBytesValue(uncachedKeys[k])
				if err == nil {
					err = unmarshalMapElement(s, creator, uncachedKeys[k], bs, c)
					if err != nil {
						return err
					}
					continue
				}
				if err != cache.ErrNotFound {
					return err
				}
			} else {
				locker.Lock()
				defer locker.Unlock()
			}
			unreloadkeys = append(unreloadkeys, uncachedKeys[k])

		}

	} else {
		unreloadkeys = uncachedKeys
	}

	if len(unreloadkeys) == 0 {
		return nil
	}
	loaded, err := loader(unreloadkeys...)
	if err != nil {
		return err
	}
	var data = make(map[string][]byte, len(loaded))
	for k := range loaded {
		v := loaded[k]
		s.Store(k, v)
		if c != nil {
			data[k], err = c.Marshal(v)
			if err != nil {
				return err
			}
		}
	}
	for k := range unreloadkeys {
		if _, ok := data[unreloadkeys[k]]; ok == false {
			data[unreloadkeys[k]] = []byte{}
		}
	}
	if c == nil {
		return nil
	}
	return c.MSetBytesValue(data, 0)
}
<|endoftext|>


<commit_before>package main

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

	_ "code.google.com/p/go-sqlite/go1/sqlite3"
	"code.google.com/p/go.crypto/bcrypt"
)

type LiteBookie struct {
	db   *sql.DB
	elog *log.Logger
	tlog *log.Logger
}

func NewLiteBookie(eWriter, tWriter io.Writer) (b *LiteBookie) {
	b = &LiteBookie{}
	b.elog = log.New(eWriter, "[Bookie] ", log.Lmicroseconds|log.Lshortfile)
	b.tlog = log.New(tWriter, "[T] ", log.Lmicroseconds)

	return b
}

func (b *LiteBookie) Open(path string) (err error) {
	b.db, err = sql.Open("sqlite3", path)
	if err != nil {
		b.elog.Println(err)
	}
	return err
}

func (b *LiteBookie) UserRegister(user, email, password string) (err error) {
	tx, err := b.db.Begin()
	if err != nil {
		b.elog.Println(err)
		return errors.New("Database failure")
	}

	sql := "INSERT INTO user(name, email, password) VALUES(?,?,?);"
	stmt, err := tx.Prepare(sql)
	if err != nil {
		b.elog.Println(err)
		return errors.New("Database failure")
	}
	defer stmt.Close()

	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		b.elog.Println(err)
		return errors.New("Temporary Registration Failure")
	}

	_, err = stmt.Exec(user, email, hash)
	if err != nil {
		b.elog.Println(err)
		return errors.New("Username or Email address not unique")
	}
	tx.Commit()

	return nil
}

func (b *LiteBookie) UserLogin(user, password string) {
}

func (b *LiteBookie) EventCreate() {
}

func (b *LiteBookie) EventOpen() {
}

func (b *LiteBookie) EventClose() {
}

func (b *LiteBookie) EventCancel() {
}

func (b *LiteBookie) EventList() {
}

func (b *LiteBookie) EventBet() {
}

func (b *LiteBookie) StreamCreate() {
}

func (b *LiteBookie) StreamList() {
}
<commit_msg>Fleshing out bookie stubs<commit_after>package main

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

	_ "code.google.com/p/go-sqlite/go1/sqlite3"
	"code.google.com/p/go.crypto/bcrypt"
)

type LiteBookie struct {
	db   *sql.DB
	elog *log.Logger
	tlog *log.Logger
}

func NewLiteBookie(eWriter, tWriter io.Writer) (b *LiteBookie) {
	b = &LiteBookie{}
	b.elog = log.New(eWriter, "[Bookie] ", log.Lmicroseconds|log.Lshortfile)
	b.tlog = log.New(tWriter, "[T] ", log.Lmicroseconds)

	return b
}

func (b *LiteBookie) sqlexec(query string, args ...interface{}) (result sql.Result, err error) {
	stmt, err := b.db.Prepare(query)
	if err != nil {
		b.elog.Println(err)
		return nil, errors.New("Temporary database issue")
	}
	defer stmt.Close()

	result, err = stmt.Exec(args...)
	if err != nil {
		b.elog.Println(err)
		return nil, errors.New("Unable to complete")
	}
	return result, err
}

func (b *LiteBookie) Open(path string) (err error) {
	b.db, err = sql.Open("sqlite3", path)
	if err != nil {
		b.elog.Println(err)
	}
	return err
}

func (b *LiteBookie) UserRegister(user, email, password string) (err error) {
	const query = "INSERT INTO user(name, email, password) VALUES(?,?,?);"

	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		b.elog.Println(err)
		return errors.New("Temporary Registration Failure")
	}

	_, err = b.sqlexec(query, user, email, hash)
	// could be a lie, fix this by recording what stage the error occurred
	if err != nil {
		return errors.New("Username or Email address not unique")
	}

	return nil
}

func (b *LiteBookie) UserLogin(user, password string) (err error) {
	const query = "SELECT password from user WHERE name=?;"
	const update = "UPDATE user SET last_login=current_timestamp WHERE name=?;"

	var hash []byte
	err = b.db.QueryRow(query, user).Scan(&hash)
	if err != nil {
		b.elog.Println(err)
		return errors.New("Invalid User")
	}

	err = bcrypt.CompareHashAndPassword(hash, []byte(password))
	if err != nil {
		b.elog.Println(err)
		return errors.New("Invalid Password")
	}

	_, err = b.sqlexec(update, user)
	if err != nil {
		return errors.New("Temporarly Unable to Login")
	}

	return err
}

func (b *LiteBookie) EventCreate() {
}

func (b *LiteBookie) EventOpen() {
}

func (b *LiteBookie) EventClose() {
}

func (b *LiteBookie) EventCancel() {
}

func (b *LiteBookie) EventList() {
}

func (b *LiteBookie) EventBet() {
}

func (b *LiteBookie) StreamCreate() {
}

func (b *LiteBookie) StreamList() {
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/brutella/hc"
	"github.com/brutella/hc/accessory"
	"github.com/brutella/hc/service"
	"github.com/tarm/serial"
	"log"
)

type MessageType byte

const (
	TypeAliveRequest      MessageType = 0x01
	TypeAliveResponse     MessageType = 0x02
	TypeSetHeightRequest  MessageType = 0x03
	TypeGetHeightRequest  MessageType = 0x04
	TypeGetHeightResponse MessageType = 0x05
	TypeStopRequest       MessageType = 0x06
	TypeGetStatusRequest  MessageType = 0x07
	TypeGetStatusResponse MessageType = 0x08
	TypeMoveUpRequest     MessageType = 0x0A
	TypeMoveDownRequest   MessageType = 0x0B
)

type Message struct {
	Type  MessageType
	Value byte
}

func receiver(c chan<- Message, p *serial.Port) {
	message := make([]byte, 3)
	buf := make([]byte, 1)

	for {
		// shift bytes to left
		message[0] = message[1]
		message[1] = message[2]

		// read new byte
		_, err := p.Read(buf)
		if err != nil {
			log.Fatal(err)
		}

		// append new byte
		message[2] = buf[0]

		// checksum
		if message[0]+message[1] != message[2] {
			continue
		}

		c <- Message{Type: MessageType(message[0]), Value: message[1]}
	}
}

func sender(c <-chan Message, p *serial.Port) {
	message := make([]byte, 3)

	for {
		// get a message
		m := <-c

		log.Println("Message to send", m)

		// fill the message buffer
		message[0] = byte(m.Type)
		message[1] = m.Value

		// calculate the checksum
		message[2] = message[0] + message[1]

		log.Println("Buffer to send", message)

		// write the buffer
		_, err := p.Write(message)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func main() {
	info := accessory.Info{
		Name:         "Desk",
		Manufacturer: "David Knezic",
	}

	acc := accessory.New(info, accessory.TypeWindow)

	service := service.NewWindow()
	service.TargetPosition.SetValue(100)

	c := &serial.Config{Name: "/dev/ttyAMA0", Baud: 9600}
	s, err := serial.OpenPort(c)
	if err != nil {
		log.Fatal(err)
	}

	var incoming chan Message = make(chan Message)
	var outgoing chan Message = make(chan Message)

	go receiver(incoming, s)
	go sender(outgoing, s)

	service.TargetPosition.OnValueRemoteUpdate(func(position int) {
		log.Println("Setting desk to", position, "percent height")

		factor := float64(position) / 100.0
		height := byte(80) + byte(100.0*factor)

		log.Println("This corresponds to", height, "cm height")

		outgoing <- Message{Type: MessageType(TypeSetHeightRequest), Value: height}

		service.CurrentPosition.SetValue(position)
	})

	acc.AddService(service.Service)

	config := hc.Config{
		Pin: "32191123",
	}

	t, err := hc.NewIPTransport(config, acc)

	if err != nil {
		log.Fatal(err)
	}

	hc.OnTermination(func() {
		t.Stop()
	})

	t.Start()
}
<commit_msg>Describe the message types, add update height event<commit_after>package main

import (
	"github.com/brutella/hc"
	"github.com/brutella/hc/accessory"
	"github.com/brutella/hc/service"
	"github.com/tarm/serial"
	"log"
)

type MessageType byte

const (
	// checking the availability of the desk
	TypeAliveRequest      MessageType = 0x01
	TypeAliveResponse     MessageType = 0x02

	// setting the height of the desk
	TypeSetHeightRequest  MessageType = 0x03

	// querying the height of the desk
	TypeGetHeightRequest  MessageType = 0x04
	TypeGetHeightResponse MessageType = 0x05

	// stopping the desk
	TypeStopRequest       MessageType = 0x06

	// TODO: to be implemented
	TypeGetStatusRequest  MessageType = 0x07
	TypeGetStatusResponse MessageType = 0x08

	// moving the desk
	TypeMoveUpRequest     MessageType = 0x0A
	TypeMoveDownRequest   MessageType = 0x0B

	// the desk notifying about a height change
	TypeUpdateHeightEvent MessageType = 0x0C
)

type Message struct {
	Type  MessageType
	Value byte
}

func receiver(c chan<- Message, p *serial.Port) {
	message := make([]byte, 3)
	buf := make([]byte, 1)

	for {
		// shift bytes to left
		message[0] = message[1]
		message[1] = message[2]

		// read new byte
		_, err := p.Read(buf)
		if err != nil {
			log.Fatal(err)
		}

		// append new byte
		message[2] = buf[0]

		// checksum
		if message[0]+message[1] != message[2] {
			continue
		}

		c <- Message{Type: MessageType(message[0]), Value: message[1]}
	}
}

func sender(c <-chan Message, p *serial.Port) {
	message := make([]byte, 3)

	for {
		// get a message
		m := <-c

		log.Println("Message to send", m)

		// fill the message buffer
		message[0] = byte(m.Type)
		message[1] = m.Value

		// calculate the checksum
		message[2] = message[0] + message[1]

		log.Println("Buffer to send", message)

		// write the buffer
		_, err := p.Write(message)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func main() {
	info := accessory.Info{
		Name:         "Desk",
		Manufacturer: "David Knezic",
	}

	acc := accessory.New(info, accessory.TypeWindow)

	service := service.NewWindow()
	service.TargetPosition.SetValue(100)

	c := &serial.Config{Name: "/dev/ttyAMA0", Baud: 9600}
	s, err := serial.OpenPort(c)
	if err != nil {
		log.Fatal(err)
	}

	var incoming chan Message = make(chan Message)
	var outgoing chan Message = make(chan Message)

	go receiver(incoming, s)
	go sender(outgoing, s)

	service.TargetPosition.OnValueRemoteUpdate(func(position int) {
		log.Println("Setting desk to", position, "percent height")

		factor := float64(position) / 100.0
		height := byte(80) + byte(100.0*factor)

		log.Println("This corresponds to", height, "cm height")

		outgoing <- Message{Type: MessageType(TypeSetHeightRequest), Value: height}

		service.CurrentPosition.SetValue(position)
	})

	acc.AddService(service.Service)

	config := hc.Config{
		Pin: "32191123",
	}

	t, err := hc.NewIPTransport(config, acc)

	if err != nil {
		log.Fatal(err)
	}

	hc.OnTermination(func() {
		t.Stop()
	})

	t.Start()
}
<|endoftext|>


<commit_before>package MQTTg

import (
	"net"
	"strconv"
	"time"
)

type Broker struct {
	Bt *Transport
	// TODO: check whether not good to use addr as key
	Clients   map[*net.UDPAddr]*ClientInfo // map[addr]*ClientInfo
	ClientIDs map[string]*ClientInfo       // map[clientID]*ClientInfo
	TopicRoot *TopicNode
}

type ClientInfo struct {
	*Client
	KeepAliveTimer *time.Timer
	Duration       time.Duration
}

func (self *ClientInfo) NewTimer() {
	self.KeepAliveTimer = time.NewTimer(self.Duration)
}

func (self *ClientInfo) ResetTimer() {
	self.KeepAliveTimer.Reset(self.Duration)
}

func (self *ClientInfo) RunTimer() {
	<-self.KeepAliveTimer.C
	self.DisconnectFromBroker()
	// TODO: logging?
}

func (self *ClientInfo) DisconnectFromBroker() {
	self.Will = nil
	self.IsConnecting = false
	self.KeepAliveTimer.Stop()
	// TODO: free used clientID due to clean session?
}

func (self *Broker) ApplyDummyClientID() string {
	return "DummyClientID:" + strconv.Itoa(len(self.ClientIDs)+1)
}

func (self *Broker) ReadLoop() error {
	for {
		m, addr, err := self.Bt.ReadMessageFrom()
		if err != nil {
			EmitError(err)
			continue
		}
		switch message := m.(type) {
		case *ConnectMessage:
			if message.Protocol.Level != MQTT_3_1_1.Level {
				// CHECK: Is false correct?
				err = self.Bt.SendMessage(NewConnackMessage(false, UnacceptableProtocolVersion), addr)
				EmitError(INVALID_PROTOCOL_LEVEL)
				continue
			}

			client, ok := self.ClientIDs[message.ClientID]
			if ok {
				// TODO: this might cause problem
				err = self.Bt.SendMessage(NewConnackMessage(false, IdentifierRejected), addr)
				EmitError(CLIENT_ID_IS_USED_ALREADY)
				continue
			}
			if len(message.ClientID) == 0 {
				message.ClientID = self.ApplyDummyClientID()
			}

			if message.Flags&Reserved_Flag == Reserved_Flag {
				// TODO: disconnect the connection
				continue
			}

			if message.Flags&UserName_Flag != UserName_Flag && message.Flags&Password_Flag == Password_Flag {
				EmitError(USERNAME_DOES_NOT_EXIST_WITH_PASSWORD)
				continue
			}
			// TODO: authorization

			client, ok = self.Clients[addr]
			sessionPresent := false
			if message.Flags&CleanSession_Flag == CleanSession_Flag || !ok {
				// TODO: need to manage QoS base processing
				duration := time.Duration(float32(message.KeepAlive) * 100000000 * 1.5)
				client = &ClientInfo{
					Client: NewClient(self.Bt, addr, message.ClientID,
						message.User, message.KeepAlive, message.Will),
					KeepAliveTimer: nil,
					Duretion:       duration,
				}
				client.NewTimer()
				self.Clients[addr] = client
				self.ClientIDs[message.ClientID] = client
			} else if message.Flags&CleanSession_Flag != CleanSession_Flag || ok {
				sessionPresent = true
			}

			if message.Flags&Will_Flag == Will_Flag {
				client.Will = message.Will
				// TODO: consider QoS and Retain as broker need
			} else {

			}

			go client.RunTimer()
			err = self.Bt.SendMessage(NewConnackMessage(sessionPresent, Accepted), addr)
		case *PublishMessage:
			if message.QoS == 3 {
				// error
				// close connection
				continue
			}

			if message.Dup {
				// re-delivered
			} else {
				// first time delivery
			}

			if message.Retain {
				// store the application message to designated topic
				if len(message.Payload) == 0 {
					// discard(remove) retained message
				}

				if message.QoS == 0 {
					// discard retained message
				}
			}

			switch message.QoS {
			// in any case, Dub must be 0
			case 0:
			case 1:
				err = self.Bt.SendMessage(NewPubackMessage(message.PacketID), addr)
			case 2:
				err = self.Bt.SendMessage(NewPubrecMessage(message.PacketID), addr)
			}
		case *PubackMessage:
			// acknowledge the sent Publish packet
		case *PubrecMessage:
			// acknowledge the sent Publish packet
			err = self.Bt.SendMessage(NewPubrelMessage(message.PacketID), addr)
		case *PubrelMessage:
			err = self.Bt.SendMessage(NewPubcompMessage(message.PacketID), addr)
		case *PubcompMessage:
			// acknowledge the sent Pubrel packet
		case *SubscribeMessage:
			// TODO: check The wild card is permitted
			codes := make([]SubscribeReturnCode, len(message.SubscribeTopics))
			for i, subTopic := range message.SubscribeTopics {
				// TODO: need to validate wheter there are same topics or not
				client, ok := self.Clients[addr]
				if ok {
					retains, code := self.TopicRoot.ApplySubscriber(client.ID, string(subTopic.Topic), subTopic.QoS)
					codes[i] = code
					client.SubTopics = append(client.SubTopics,
						SubscribeTopic{SubscribeAck,
							subTopic.Topic,
							uint8(code),
						})
					if len(retains) > 0 {
						for k, v := range retains {
							//Publish(k,v)
						}
					}
				} else {
					codes[i] = SubscribeFailure // TODO: correct?
				}

			}
			err = self.Bt.SendMessage(NewSubackMessage(message.PacketID, codes), addr)
		case *UnsubscribeMessage:
			client, ok := self.Clients[addr]
			if !ok {
				EmitError(CLIENT_NOT_EXIST)
				continue // TODO: ?
			}
			if len(message.TopicNames) == 0 {
				// protocol violation
			}
			// TODO: optimize here
			result := []SubscribeTopic{}
			for _, t := range client.SubTopics {
				del := false
				for _, name := range message.TopicNames {
					if string(t.Topic) == string(name) {
						del = true
					}
				}
				if !del {
					result = append(result, t)
				}
			}
			client.SubTopics = result
			err = self.Bt.SendMessage(NewUnsubackMessage(message.PacketID), addr)
		case *PingreqMessage:
			// Pingresp
			// TODO: calc elapsed time from previous pingreq.
			//       and store the time to duration of Transport
			err = self.Bt.SendMessage(NewPingrespMessage(), addr)
			client, ok := self.Clients[addr]
			if !ok {
				EmitError(CLIENT_NOT_EXIST)
				continue
			}
			client.ResetTimer()
			go client.RunTimer()
		case *DisconnectMessage:
			client, ok := self.Clients[addr]
			if !ok {
				EmitError(CLIENT_NOT_EXIST)
				continue
			}
			client.DisconnectFromBroker()

			// close the client
		default:
			// when invalid messages come
			err = INVALID_MESSAGE_CAME
		}
		if err != nil {
			EmitError(err)
		}
	}
}
<commit_msg>added generator instead of NewTimer<commit_after>package MQTTg

import (
	"net"
	"strconv"
	"time"
)

type Broker struct {
	Bt *Transport
	// TODO: check whether not good to use addr as key
	Clients   map[*net.UDPAddr]*ClientInfo // map[addr]*ClientInfo
	ClientIDs map[string]*ClientInfo       // map[clientID]*ClientInfo
	TopicRoot *TopicNode
}

type ClientInfo struct {
	*Client
	KeepAliveTimer *time.Timer
	Duration       time.Duration
}

func NewClientInfo(c *Client, duration time.Duration) *ClientInfo {
	return &ClientInfo{
		Client:         c,
		KeepAliveTimer: time.NewTimer(duration),
		Duration:       duration,
	}
}

func (self *ClientInfo) ResetTimer() {
	self.KeepAliveTimer.Reset(self.Duration)
}

func (self *ClientInfo) RunTimer() {
	<-self.KeepAliveTimer.C
	self.DisconnectFromBroker()
	// TODO: logging?
}

func (self *ClientInfo) DisconnectFromBroker() {
	self.Will = nil
	self.IsConnecting = false
	self.KeepAliveTimer.Stop()
	// TODO: free used clientID due to clean session?
}

func (self *Broker) ApplyDummyClientID() string {
	return "DummyClientID:" + strconv.Itoa(len(self.ClientIDs)+1)
}

func (self *Broker) ReadLoop() error {
	for {
		m, addr, err := self.Bt.ReadMessageFrom()
		if err != nil {
			EmitError(err)
			continue
		}
		switch message := m.(type) {
		case *ConnectMessage:
			if message.Protocol.Level != MQTT_3_1_1.Level {
				// CHECK: Is false correct?
				err = self.Bt.SendMessage(NewConnackMessage(false, UnacceptableProtocolVersion), addr)
				EmitError(INVALID_PROTOCOL_LEVEL)
				continue
			}

			client, ok := self.ClientIDs[message.ClientID]
			if ok {
				// TODO: this might cause problem
				err = self.Bt.SendMessage(NewConnackMessage(false, IdentifierRejected), addr)
				EmitError(CLIENT_ID_IS_USED_ALREADY)
				continue
			}
			if len(message.ClientID) == 0 {
				message.ClientID = self.ApplyDummyClientID()
			}

			if message.Flags&Reserved_Flag == Reserved_Flag {
				// TODO: disconnect the connection
				continue
			}

			if message.Flags&UserName_Flag != UserName_Flag && message.Flags&Password_Flag == Password_Flag {
				EmitError(USERNAME_DOES_NOT_EXIST_WITH_PASSWORD)
				continue
			}
			// TODO: authorization

			client, ok = self.Clients[addr]
			sessionPresent := false
			if message.Flags&CleanSession_Flag == CleanSession_Flag || !ok {
				// TODO: need to manage QoS base processing
				duration := time.Duration(float32(message.KeepAlive) * 100000000 * 1.5)
				client = NewClientInfo(NewClient(self.Bt, addr, message.ClientID,
					message.User, message.KeepAlive, message.Will), duration)

				self.Clients[addr] = client
				self.ClientIDs[message.ClientID] = client
			} else if message.Flags&CleanSession_Flag != CleanSession_Flag || ok {
				sessionPresent = true
			}

			if message.Flags&Will_Flag == Will_Flag {
				client.Will = message.Will
				// TODO: consider QoS and Retain as broker need
			} else {

			}

			go client.RunTimer()
			err = self.Bt.SendMessage(NewConnackMessage(sessionPresent, Accepted), addr)
		case *PublishMessage:
			if message.QoS == 3 {
				// error
				// close connection
				continue
			}

			if message.Dup {
				// re-delivered
			} else {
				// first time delivery
			}

			if message.Retain {
				// store the application message to designated topic
				if len(message.Payload) == 0 {
					// discard(remove) retained message
				}

				if message.QoS == 0 {
					// discard retained message
				}
			}

			switch message.QoS {
			// in any case, Dub must be 0
			case 0:
			case 1:
				err = self.Bt.SendMessage(NewPubackMessage(message.PacketID), addr)
			case 2:
				err = self.Bt.SendMessage(NewPubrecMessage(message.PacketID), addr)
			}
		case *PubackMessage:
			// acknowledge the sent Publish packet
		case *PubrecMessage:
			// acknowledge the sent Publish packet
			err = self.Bt.SendMessage(NewPubrelMessage(message.PacketID), addr)
		case *PubrelMessage:
			err = self.Bt.SendMessage(NewPubcompMessage(message.PacketID), addr)
		case *PubcompMessage:
			// acknowledge the sent Pubrel packet
		case *SubscribeMessage:
			// TODO: check The wild card is permitted
			codes := make([]SubscribeReturnCode, len(message.SubscribeTopics))
			for i, subTopic := range message.SubscribeTopics {
				// TODO: need to validate wheter there are same topics or not
				client, ok := self.Clients[addr]
				if ok {
					retains, code := self.TopicRoot.ApplySubscriber(client.ID, string(subTopic.Topic), subTopic.QoS)
					codes[i] = code
					client.SubTopics = append(client.SubTopics,
						SubscribeTopic{SubscribeAck,
							subTopic.Topic,
							uint8(code),
						})
					if len(retains) > 0 {
						for k, v := range retains {
							//Publish(k,v)
						}
					}
				} else {
					codes[i] = SubscribeFailure // TODO: correct?
				}

			}
			err = self.Bt.SendMessage(NewSubackMessage(message.PacketID, codes), addr)
		case *UnsubscribeMessage:
			client, ok := self.Clients[addr]
			if !ok {
				EmitError(CLIENT_NOT_EXIST)
				continue // TODO: ?
			}
			if len(message.TopicNames) == 0 {
				// protocol violation
			}
			// TODO: optimize here
			result := []SubscribeTopic{}
			for _, t := range client.SubTopics {
				del := false
				for _, name := range message.TopicNames {
					if string(t.Topic) == string(name) {
						del = true
					}
				}
				if !del {
					result = append(result, t)
				}
			}
			client.SubTopics = result
			err = self.Bt.SendMessage(NewUnsubackMessage(message.PacketID), addr)
		case *PingreqMessage:
			// Pingresp
			// TODO: calc elapsed time from previous pingreq.
			//       and store the time to duration of Transport
			err = self.Bt.SendMessage(NewPingrespMessage(), addr)
			client, ok := self.Clients[addr]
			if !ok {
				EmitError(CLIENT_NOT_EXIST)
				continue
			}
			client.ResetTimer()
			go client.RunTimer()
		case *DisconnectMessage:
			client, ok := self.Clients[addr]
			if !ok {
				EmitError(CLIENT_NOT_EXIST)
				continue
			}
			client.DisconnectFromBroker()

			// close the client
		default:
			// when invalid messages come
			err = INVALID_MESSAGE_CAME
		}
		if err != nil {
			EmitError(err)
		}
	}
}
<|endoftext|>