<commit_before>package main

import (
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"gopkg.in/yaml.v2"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/ioprogress"
	"github.com/lxc/lxd/shared/logger"
	"github.com/lxc/lxd/shared/version"

	log "gopkg.in/inconshreveable/log15.v2"
)

// Simplestream cache
type imageStreamCacheEntry struct {
	Aliases      []api.ImageAliasesEntry `yaml:"aliases"`
	Certificate  string                  `yaml:"certificate"`
	Fingerprints []string                `yaml:"fingerprints"`

	expiry time.Time
	remote lxd.ImageServer
}

var imageStreamCache = map[string]*imageStreamCacheEntry{}
var imageStreamCacheLock sync.Mutex

var imagesDownloading = map[string]chan bool{}
var imagesDownloadingLock sync.Mutex

func imageSaveStreamCache() error {
	data, err := yaml.Marshal(&imageStreamCache)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(shared.CachePath("simplestreams.yaml"), data, 0600)
	if err != nil {
		return err
	}

	return nil
}

func imageLoadStreamCache(d *Daemon) error {
	imageStreamCacheLock.Lock()
	defer imageStreamCacheLock.Unlock()

	if !shared.PathExists(shared.CachePath("simplestreams.yaml")) {
		return nil
	}

	content, err := ioutil.ReadFile(shared.CachePath("simplestreams.yaml"))
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(content, imageStreamCache)
	if err != nil {
		return err
	}

	for url, entry := range imageStreamCache {
		if entry.remote == nil {
			remote, err := lxd.ConnectSimpleStreams(url, &lxd.ConnectionArgs{
				TLSServerCert: entry.Certificate,
				UserAgent:     version.UserAgent,
				Proxy:         d.proxy,
			})
			if err != nil {
				continue
			}

			entry.remote = remote
		}
	}

	return nil
}

// ImageDownload resolves the image fingerprint and if not in the database, downloads it
func (d *Daemon) ImageDownload(op *operation, server string, protocol string, certificate string, secret string, alias string, forContainer bool, autoUpdate bool, storagePool string, preferCached bool) (*api.Image, error) {
	var err error
	var ctxMap log.Ctx

	var remote lxd.ImageServer
	var info *api.Image

	// Default protocol is LXD
	if protocol == "" {
		protocol = "lxd"
	}

	// Default the fingerprint to the alias string we received
	fp := alias

	// Attempt to resolve the alias
	if protocol == "simplestreams" {
		imageStreamCacheLock.Lock()
		entry, _ := imageStreamCache[server]
		if entry == nil || entry.expiry.Before(time.Now()) {
			// Add a new entry to the cache
			refresh := func() (*imageStreamCacheEntry, error) {
				// Setup simplestreams client
				remote, err = lxd.ConnectSimpleStreams(server, &lxd.ConnectionArgs{
					TLSServerCert: certificate,
					UserAgent:     version.UserAgent,
					Proxy:         d.proxy,
				})
				if err != nil {
					return nil, err
				}

				// Get all aliases
				aliases, err := remote.GetImageAliases()
				if err != nil {
					return nil, err
				}

				// Get all fingerprints
				images, err := remote.GetImages()
				if err != nil {
					return nil, err
				}

				fingerprints := []string{}
				for _, image := range images {
					fingerprints = append(fingerprints, image.Fingerprint)
				}

				// Generate cache entry
				entry = &imageStreamCacheEntry{remote: remote, Aliases: aliases, Certificate: certificate, Fingerprints: fingerprints, expiry: time.Now().Add(time.Hour)}
				imageStreamCache[server] = entry
				imageSaveStreamCache()

				return entry, nil
			}

			newEntry, err := refresh()
			if err == nil {
				// Cache refreshed
				entry = newEntry
			} else if entry != nil {
				// Failed to fetch entry but existing cache
				logger.Warn("Unable to refresh cache, using stale entry", log.Ctx{"server": server})
				entry.expiry = time.Now().Add(time.Hour)
			} else {
				// Failed to fetch entry and nothing in cache
				imageStreamCacheLock.Unlock()
				return nil, err
			}
		} else {
			// use the existing entry
			logger.Debug("Using SimpleStreams cache entry", log.Ctx{"server": server, "expiry": entry.expiry})
			remote = entry.remote
		}
		imageStreamCacheLock.Unlock()

		// Look for a matching alias
		for _, entry := range entry.Aliases {
			if entry.Name != fp {
				continue
			}

			fp = entry.Target
			break
		}

		// Expand partial fingerprints
		matches := []string{}
		for _, entry := range entry.Fingerprints {
			if strings.HasPrefix(entry, fp) {
				matches = append(matches, entry)
			}
		}

		if len(matches) == 1 {
			fp = matches[0]
		} else if len(matches) > 1 {
			return nil, fmt.Errorf("Provided partial image fingerprint matches more than one image")
		}
	} else if protocol == "lxd" {
		// Setup LXD client
		remote, err = lxd.ConnectPublicLXD(server, &lxd.ConnectionArgs{
			TLSServerCert: certificate,
			UserAgent:     version.UserAgent,
			Proxy:         d.proxy,
		})
		if err != nil {
			return nil, err
		}

		// For public images, handle aliases and initial metadata
		if secret == "" {
			// Look for a matching alias
			entry, _, err := remote.GetImageAlias(fp)
			if err == nil {
				fp = entry.Target
			}

			// Expand partial fingerprints
			info, _, err = remote.GetImage(fp)
			if err != nil {
				return nil, err
			}

			fp = info.Fingerprint
		}
	}

	// If auto-update is on and we're being given the image by
	// alias, try to use a locally cached image matching the given
	// server/protocol/alias, regardless of whether it's stale or
	// not (we can assume that it will be not *too* stale since
	// auto-update is on).
	interval := daemonConfig["images.auto_update_interval"].GetInt64()
	if preferCached && interval > 0 && alias != fp {
		cachedFingerprint, err := dbImageSourceGetCachedFingerprint(d.db, server, protocol, alias)
		if err == nil && cachedFingerprint != fp {
			fp = cachedFingerprint
		}
	}

	// Check if the image already exists (partial hash match)
	_, imgInfo, err := dbImageGet(d.db, fp, false, true)
	if err == nil {
		logger.Debug("Image already exists in the db", log.Ctx{"image": fp})
		info = imgInfo

		// If not requested in a particular pool, we're done.
		if storagePool == "" {
			return info, nil
		}

		// Get the ID of the target storage pool
		poolID, err := dbStoragePoolGetID(d.db, storagePool)
		if err != nil {
			return nil, err
		}

		// Check if the image is already in the pool
		poolIDs, err := dbImageGetPools(d.db, info.Fingerprint)
		if err != nil {
			return nil, err
		}

		if shared.Int64InSlice(poolID, poolIDs) {
			logger.Debugf("Image already exists on storage pool \"%s\".", storagePool)
			return info, nil
		}

		// Import the image in the pool
		logger.Debugf("Image does not exist on storage pool \"%s\".", storagePool)

		err = imageCreateInPool(d, info, storagePool)
		if err != nil {
			logger.Debugf("Failed to create image on storage pool \"%s\": %s.", storagePool, err)
			return nil, err
		}

		logger.Debugf("Created image on storage pool \"%s\".", storagePool)
		return info, nil
	}

	// Deal with parallel downloads
	imagesDownloadingLock.Lock()
	if waitChannel, ok := imagesDownloading[fp]; ok {
		// We are already downloading the image
		imagesDownloadingLock.Unlock()

		logger.Debug(
			"Already downloading the image, waiting for it to succeed",
			log.Ctx{"image": fp})

		// Wait until the download finishes (channel closes)
		<-waitChannel

		// Grab the database entry
		_, imgInfo, err := dbImageGet(d.db, fp, false, true)
		if err != nil {
			// Other download failed, lets try again
			logger.Error("Other image download didn't succeed", log.Ctx{"image": fp})
		} else {
			// Other download succeeded, we're done
			return imgInfo, nil
		}
	} else {
		imagesDownloadingLock.Unlock()
	}

	// Add the download to the queue
	imagesDownloadingLock.Lock()
	imagesDownloading[fp] = make(chan bool)
	imagesDownloadingLock.Unlock()

	// Unlock once this func ends.
	defer func() {
		imagesDownloadingLock.Lock()
		if waitChannel, ok := imagesDownloading[fp]; ok {
			close(waitChannel)
			delete(imagesDownloading, fp)
		}
		imagesDownloadingLock.Unlock()
	}()

	// Begin downloading
	if op == nil {
		ctxMap = log.Ctx{"alias": alias, "server": server}
	} else {
		ctxMap = log.Ctx{"trigger": op.url, "image": fp, "operation": op.id, "alias": alias, "server": server}
	}
	logger.Info("Downloading image", ctxMap)

	// Cleanup any leftover from a past attempt
	destDir := shared.VarPath("images")
	destName := filepath.Join(destDir, fp)

	failure := true
	cleanup := func() {
		if failure {
			os.Remove(destName)
			os.Remove(destName + ".rootfs")
		}
	}
	defer cleanup()

	// Setup a progress handler
	progress := func(progress lxd.ProgressData) {
		if op == nil {
			return
		}

		meta := op.metadata
		if meta == nil {
			meta = make(map[string]interface{})
		}

		if meta["download_progress"] != progress.Text {
			meta["download_progress"] = progress.Text
			op.UpdateMetadata(meta)
		}
	}

	if protocol == "lxd" || protocol == "simplestreams" {
		// Create the target files
		dest, err := os.Create(destName)
		if err != nil {
			return nil, err
		}
		defer dest.Close()

		destRootfs, err := os.Create(destName + ".rootfs")
		if err != nil {
			return nil, err
		}
		defer destRootfs.Close()

		// Get the image information
		if info == nil {
			if secret != "" {
				info, _, err = remote.GetPrivateImage(fp, secret)
			} else {
				info, _, err = remote.GetImage(fp)
			}
			if err != nil {
				return nil, err
			}
		}

		// Download the image
		var resp *lxd.ImageFileResponse
		request := lxd.ImageFileRequest{
			MetaFile:        io.WriteSeeker(dest),
			RootfsFile:      io.WriteSeeker(destRootfs),
			ProgressHandler: progress,
		}

		if secret != "" {
			resp, err = remote.GetPrivateImageFile(fp, secret, request)
		} else {
			resp, err = remote.GetImageFile(fp, request)
		}
		if err != nil {
			return nil, err
		}

		// Deal with unified images
		if resp.RootfsSize == 0 {
			err := os.Remove(destName + ".rootfs")
			if err != nil {
				return nil, err
			}
		}
	} else if protocol == "direct" {
		// Setup HTTP client
		httpClient, err := d.httpClient(certificate)
		if err != nil {
			return nil, err
		}

		req, err := http.NewRequest("GET", server, nil)
		if err != nil {
			return nil, err
		}

		req.Header.Set("User-Agent", version.UserAgent)

		// Make the request
		raw, err := httpClient.Do(req)
		if err != nil {
			return nil, err
		}

		if raw.StatusCode != http.StatusOK {
			return nil, fmt.Errorf("Unable to fetch %s: %s", server, raw.Status)
		}

		// Progress handler
		body := &ioprogress.ProgressReader{
			ReadCloser: raw.Body,
			Tracker: &ioprogress.ProgressTracker{
				Length: raw.ContentLength,
				Handler: func(percent int64, speed int64) {
					progress(lxd.ProgressData{Text: fmt.Sprintf("%d%% (%s/s)", percent, shared.GetByteSizeString(speed, 2))})
				},
			},
		}

		// Create the target files
		f, err := os.Create(destName)
		if err != nil {
			return nil, err
		}
		defer f.Close()

		// Hashing
		sha256 := sha256.New()

		// Download the image
		size, err := io.Copy(io.MultiWriter(f, sha256), body)
		if err != nil {
			return nil, err
		}

		// Validate hash
		result := fmt.Sprintf("%x", sha256.Sum(nil))
		if result != fp {
			return nil, fmt.Errorf("Hash mismatch for %s: %s != %s", server, result, fp)
		}

		// Parse the image
		imageMeta, err := getImageMetadata(destName)
		if err != nil {
			return nil, err
		}

		info.Size = size
		info.Architecture = imageMeta.Architecture
		info.CreatedAt = time.Unix(imageMeta.CreationDate, 0)
		info.ExpiresAt = time.Unix(imageMeta.ExpiryDate, 0)
		info.Properties = imageMeta.Properties
	}

	// Override visiblity
	info.Public = false

	// We want to enable auto-update only if we were passed an
	// alias name, so we can figure when the associated
	// fingerprint changes in the remote.
	if alias != fp {
		info.AutoUpdate = autoUpdate
	}

	// Create the database entry
	err = dbImageInsert(d.db, info.Fingerprint, info.Filename, info.Size, info.Public, info.AutoUpdate, info.Architecture, info.CreatedAt, info.ExpiresAt, info.Properties)
	if err != nil {
		return nil, err
	}

	// Image is in the DB now, don't wipe on-disk files on failure
	failure = false

	if alias != fp {
		id, _, err := dbImageGet(d.db, fp, false, true)
		if err != nil {
			return nil, err
		}

		err = dbImageSourceInsert(d.db, id, server, protocol, certificate, alias)
		if err != nil {
			return nil, err
		}
	}

	// Import into the requested storage pool
	if storagePool != "" {
		err = imageCreateInPool(d, info, storagePool)
		if err != nil {
			return nil, err
		}
	}

	// Mark the image as "cached" if downloading for a container
	if forContainer {
		err := dbImageLastAccessInit(d.db, fp)
		if err != nil {
			return nil, err
		}
	}

	logger.Info("Image downloaded", ctxMap)
	return info, nil
}
<commit_msg>lxd/images: Initialize image info in direct case<commit_after>package main

import (
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"gopkg.in/yaml.v2"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/ioprogress"
	"github.com/lxc/lxd/shared/logger"
	"github.com/lxc/lxd/shared/version"

	log "gopkg.in/inconshreveable/log15.v2"
)

// Simplestream cache
type imageStreamCacheEntry struct {
	Aliases      []api.ImageAliasesEntry `yaml:"aliases"`
	Certificate  string                  `yaml:"certificate"`
	Fingerprints []string                `yaml:"fingerprints"`

	expiry time.Time
	remote lxd.ImageServer
}

var imageStreamCache = map[string]*imageStreamCacheEntry{}
var imageStreamCacheLock sync.Mutex

var imagesDownloading = map[string]chan bool{}
var imagesDownloadingLock sync.Mutex

func imageSaveStreamCache() error {
	data, err := yaml.Marshal(&imageStreamCache)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(shared.CachePath("simplestreams.yaml"), data, 0600)
	if err != nil {
		return err
	}

	return nil
}

func imageLoadStreamCache(d *Daemon) error {
	imageStreamCacheLock.Lock()
	defer imageStreamCacheLock.Unlock()

	if !shared.PathExists(shared.CachePath("simplestreams.yaml")) {
		return nil
	}

	content, err := ioutil.ReadFile(shared.CachePath("simplestreams.yaml"))
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(content, imageStreamCache)
	if err != nil {
		return err
	}

	for url, entry := range imageStreamCache {
		if entry.remote == nil {
			remote, err := lxd.ConnectSimpleStreams(url, &lxd.ConnectionArgs{
				TLSServerCert: entry.Certificate,
				UserAgent:     version.UserAgent,
				Proxy:         d.proxy,
			})
			if err != nil {
				continue
			}

			entry.remote = remote
		}
	}

	return nil
}

// ImageDownload resolves the image fingerprint and if not in the database, downloads it
func (d *Daemon) ImageDownload(op *operation, server string, protocol string, certificate string, secret string, alias string, forContainer bool, autoUpdate bool, storagePool string, preferCached bool) (*api.Image, error) {
	var err error
	var ctxMap log.Ctx

	var remote lxd.ImageServer
	var info *api.Image

	// Default protocol is LXD
	if protocol == "" {
		protocol = "lxd"
	}

	// Default the fingerprint to the alias string we received
	fp := alias

	// Attempt to resolve the alias
	if protocol == "simplestreams" {
		imageStreamCacheLock.Lock()
		entry, _ := imageStreamCache[server]
		if entry == nil || entry.expiry.Before(time.Now()) {
			// Add a new entry to the cache
			refresh := func() (*imageStreamCacheEntry, error) {
				// Setup simplestreams client
				remote, err = lxd.ConnectSimpleStreams(server, &lxd.ConnectionArgs{
					TLSServerCert: certificate,
					UserAgent:     version.UserAgent,
					Proxy:         d.proxy,
				})
				if err != nil {
					return nil, err
				}

				// Get all aliases
				aliases, err := remote.GetImageAliases()
				if err != nil {
					return nil, err
				}

				// Get all fingerprints
				images, err := remote.GetImages()
				if err != nil {
					return nil, err
				}

				fingerprints := []string{}
				for _, image := range images {
					fingerprints = append(fingerprints, image.Fingerprint)
				}

				// Generate cache entry
				entry = &imageStreamCacheEntry{remote: remote, Aliases: aliases, Certificate: certificate, Fingerprints: fingerprints, expiry: time.Now().Add(time.Hour)}
				imageStreamCache[server] = entry
				imageSaveStreamCache()

				return entry, nil
			}

			newEntry, err := refresh()
			if err == nil {
				// Cache refreshed
				entry = newEntry
			} else if entry != nil {
				// Failed to fetch entry but existing cache
				logger.Warn("Unable to refresh cache, using stale entry", log.Ctx{"server": server})
				entry.expiry = time.Now().Add(time.Hour)
			} else {
				// Failed to fetch entry and nothing in cache
				imageStreamCacheLock.Unlock()
				return nil, err
			}
		} else {
			// use the existing entry
			logger.Debug("Using SimpleStreams cache entry", log.Ctx{"server": server, "expiry": entry.expiry})
			remote = entry.remote
		}
		imageStreamCacheLock.Unlock()

		// Look for a matching alias
		for _, entry := range entry.Aliases {
			if entry.Name != fp {
				continue
			}

			fp = entry.Target
			break
		}

		// Expand partial fingerprints
		matches := []string{}
		for _, entry := range entry.Fingerprints {
			if strings.HasPrefix(entry, fp) {
				matches = append(matches, entry)
			}
		}

		if len(matches) == 1 {
			fp = matches[0]
		} else if len(matches) > 1 {
			return nil, fmt.Errorf("Provided partial image fingerprint matches more than one image")
		}
	} else if protocol == "lxd" {
		// Setup LXD client
		remote, err = lxd.ConnectPublicLXD(server, &lxd.ConnectionArgs{
			TLSServerCert: certificate,
			UserAgent:     version.UserAgent,
			Proxy:         d.proxy,
		})
		if err != nil {
			return nil, err
		}

		// For public images, handle aliases and initial metadata
		if secret == "" {
			// Look for a matching alias
			entry, _, err := remote.GetImageAlias(fp)
			if err == nil {
				fp = entry.Target
			}

			// Expand partial fingerprints
			info, _, err = remote.GetImage(fp)
			if err != nil {
				return nil, err
			}

			fp = info.Fingerprint
		}
	}

	// If auto-update is on and we're being given the image by
	// alias, try to use a locally cached image matching the given
	// server/protocol/alias, regardless of whether it's stale or
	// not (we can assume that it will be not *too* stale since
	// auto-update is on).
	interval := daemonConfig["images.auto_update_interval"].GetInt64()
	if preferCached && interval > 0 && alias != fp {
		cachedFingerprint, err := dbImageSourceGetCachedFingerprint(d.db, server, protocol, alias)
		if err == nil && cachedFingerprint != fp {
			fp = cachedFingerprint
		}
	}

	// Check if the image already exists (partial hash match)
	_, imgInfo, err := dbImageGet(d.db, fp, false, true)
	if err == nil {
		logger.Debug("Image already exists in the db", log.Ctx{"image": fp})
		info = imgInfo

		// If not requested in a particular pool, we're done.
		if storagePool == "" {
			return info, nil
		}

		// Get the ID of the target storage pool
		poolID, err := dbStoragePoolGetID(d.db, storagePool)
		if err != nil {
			return nil, err
		}

		// Check if the image is already in the pool
		poolIDs, err := dbImageGetPools(d.db, info.Fingerprint)
		if err != nil {
			return nil, err
		}

		if shared.Int64InSlice(poolID, poolIDs) {
			logger.Debugf("Image already exists on storage pool \"%s\".", storagePool)
			return info, nil
		}

		// Import the image in the pool
		logger.Debugf("Image does not exist on storage pool \"%s\".", storagePool)

		err = imageCreateInPool(d, info, storagePool)
		if err != nil {
			logger.Debugf("Failed to create image on storage pool \"%s\": %s.", storagePool, err)
			return nil, err
		}

		logger.Debugf("Created image on storage pool \"%s\".", storagePool)
		return info, nil
	}

	// Deal with parallel downloads
	imagesDownloadingLock.Lock()
	if waitChannel, ok := imagesDownloading[fp]; ok {
		// We are already downloading the image
		imagesDownloadingLock.Unlock()

		logger.Debug(
			"Already downloading the image, waiting for it to succeed",
			log.Ctx{"image": fp})

		// Wait until the download finishes (channel closes)
		<-waitChannel

		// Grab the database entry
		_, imgInfo, err := dbImageGet(d.db, fp, false, true)
		if err != nil {
			// Other download failed, lets try again
			logger.Error("Other image download didn't succeed", log.Ctx{"image": fp})
		} else {
			// Other download succeeded, we're done
			return imgInfo, nil
		}
	} else {
		imagesDownloadingLock.Unlock()
	}

	// Add the download to the queue
	imagesDownloadingLock.Lock()
	imagesDownloading[fp] = make(chan bool)
	imagesDownloadingLock.Unlock()

	// Unlock once this func ends.
	defer func() {
		imagesDownloadingLock.Lock()
		if waitChannel, ok := imagesDownloading[fp]; ok {
			close(waitChannel)
			delete(imagesDownloading, fp)
		}
		imagesDownloadingLock.Unlock()
	}()

	// Begin downloading
	if op == nil {
		ctxMap = log.Ctx{"alias": alias, "server": server}
	} else {
		ctxMap = log.Ctx{"trigger": op.url, "image": fp, "operation": op.id, "alias": alias, "server": server}
	}
	logger.Info("Downloading image", ctxMap)

	// Cleanup any leftover from a past attempt
	destDir := shared.VarPath("images")
	destName := filepath.Join(destDir, fp)

	failure := true
	cleanup := func() {
		if failure {
			os.Remove(destName)
			os.Remove(destName + ".rootfs")
		}
	}
	defer cleanup()

	// Setup a progress handler
	progress := func(progress lxd.ProgressData) {
		if op == nil {
			return
		}

		meta := op.metadata
		if meta == nil {
			meta = make(map[string]interface{})
		}

		if meta["download_progress"] != progress.Text {
			meta["download_progress"] = progress.Text
			op.UpdateMetadata(meta)
		}
	}

	if protocol == "lxd" || protocol == "simplestreams" {
		// Create the target files
		dest, err := os.Create(destName)
		if err != nil {
			return nil, err
		}
		defer dest.Close()

		destRootfs, err := os.Create(destName + ".rootfs")
		if err != nil {
			return nil, err
		}
		defer destRootfs.Close()

		// Get the image information
		if info == nil {
			if secret != "" {
				info, _, err = remote.GetPrivateImage(fp, secret)
			} else {
				info, _, err = remote.GetImage(fp)
			}
			if err != nil {
				return nil, err
			}
		}

		// Download the image
		var resp *lxd.ImageFileResponse
		request := lxd.ImageFileRequest{
			MetaFile:        io.WriteSeeker(dest),
			RootfsFile:      io.WriteSeeker(destRootfs),
			ProgressHandler: progress,
		}

		if secret != "" {
			resp, err = remote.GetPrivateImageFile(fp, secret, request)
		} else {
			resp, err = remote.GetImageFile(fp, request)
		}
		if err != nil {
			return nil, err
		}

		// Deal with unified images
		if resp.RootfsSize == 0 {
			err := os.Remove(destName + ".rootfs")
			if err != nil {
				return nil, err
			}
		}
	} else if protocol == "direct" {
		// Setup HTTP client
		httpClient, err := d.httpClient(certificate)
		if err != nil {
			return nil, err
		}

		req, err := http.NewRequest("GET", server, nil)
		if err != nil {
			return nil, err
		}

		req.Header.Set("User-Agent", version.UserAgent)

		// Make the request
		raw, err := httpClient.Do(req)
		if err != nil {
			return nil, err
		}

		if raw.StatusCode != http.StatusOK {
			return nil, fmt.Errorf("Unable to fetch %s: %s", server, raw.Status)
		}

		// Progress handler
		body := &ioprogress.ProgressReader{
			ReadCloser: raw.Body,
			Tracker: &ioprogress.ProgressTracker{
				Length: raw.ContentLength,
				Handler: func(percent int64, speed int64) {
					progress(lxd.ProgressData{Text: fmt.Sprintf("%d%% (%s/s)", percent, shared.GetByteSizeString(speed, 2))})
				},
			},
		}

		// Create the target files
		f, err := os.Create(destName)
		if err != nil {
			return nil, err
		}
		defer f.Close()

		// Hashing
		sha256 := sha256.New()

		// Download the image
		size, err := io.Copy(io.MultiWriter(f, sha256), body)
		if err != nil {
			return nil, err
		}

		// Validate hash
		result := fmt.Sprintf("%x", sha256.Sum(nil))
		if result != fp {
			return nil, fmt.Errorf("Hash mismatch for %s: %s != %s", server, result, fp)
		}

		// Parse the image
		imageMeta, err := getImageMetadata(destName)
		if err != nil {
			return nil, err
		}

		info = &api.Image{}
		info.Fingerprint = fp
		info.Size = size
		info.Architecture = imageMeta.Architecture
		info.CreatedAt = time.Unix(imageMeta.CreationDate, 0)
		info.ExpiresAt = time.Unix(imageMeta.ExpiryDate, 0)
		info.Properties = imageMeta.Properties
	}

	// Override visiblity
	info.Public = false

	// We want to enable auto-update only if we were passed an
	// alias name, so we can figure when the associated
	// fingerprint changes in the remote.
	if alias != fp {
		info.AutoUpdate = autoUpdate
	}

	// Create the database entry
	err = dbImageInsert(d.db, info.Fingerprint, info.Filename, info.Size, info.Public, info.AutoUpdate, info.Architecture, info.CreatedAt, info.ExpiresAt, info.Properties)
	if err != nil {
		return nil, err
	}

	// Image is in the DB now, don't wipe on-disk files on failure
	failure = false

	if alias != fp {
		id, _, err := dbImageGet(d.db, fp, false, true)
		if err != nil {
			return nil, err
		}

		err = dbImageSourceInsert(d.db, id, server, protocol, certificate, alias)
		if err != nil {
			return nil, err
		}
	}

	// Import into the requested storage pool
	if storagePool != "" {
		err = imageCreateInPool(d, info, storagePool)
		if err != nil {
			return nil, err
		}
	}

	// Mark the image as "cached" if downloading for a container
	if forContainer {
		err := dbImageLastAccessInit(d.db, fp)
		if err != nil {
			return nil, err
		}
	}

	logger.Info("Image downloaded", ctxMap)
	return info, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/url"
	"strconv"
	"time"

	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
)

type cmdShutdown struct {
	global *cmdGlobal

	flagForce   bool
	flagTimeout int
}

func (c *cmdShutdown) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = "shutdown"
	cmd.Short = "Tell LXD to shutdown all containers and exit"
	cmd.Long = `Description:
  Tell LXD to shutdown all containers and exit

  This will tell LXD to start a clean shutdown of all containers,
  followed by having itself shutdown and exit.

  This can take quite a while as containers can take a long time to
  shutdown, especially if a non-standard timeout was configured for them.
`
	cmd.RunE = c.Run
	cmd.Flags().IntVarP(&c.flagTimeout, "timeout", "t", 0, "Number of seconds to wait before giving up"+"``")
	cmd.Flags().BoolVarP(&c.flagForce, "force", "f", false, "Force shutdown instead of waiting for running operations to finish"+"``")

	return cmd
}

func (c *cmdShutdown) Run(cmd *cobra.Command, args []string) error {
	connArgs := &lxd.ConnectionArgs{
		SkipGetServer: true,
	}

	d, err := lxd.ConnectLXDUnix("", connArgs)
	if err != nil {
		return err
	}

	v := url.Values{}
	v.Set("force", strconv.FormatBool(c.flagForce))

	chResult := make(chan error, 1)
	go func() {
		defer close(chResult)

		// Request shutdown, this shouldn't return until daemon has stopped.
		_, _, err = d.RawQuery("PUT", fmt.Sprintf("/internal/shutdown?%s", v.Encode()), nil, "")
		if err != nil {
			chResult <- err
			return
		}

		// Try connecting to events endpoint to check the daemon has really shutdown.
		monitor, err := d.GetEvents()
		if err != nil {
			return // Daemon has stopped.
		}

		monitor.Wait() // Wait for daemon to stop.
	}()

	if c.flagTimeout > 0 {
		select {
		case err = <-chResult:
			return err
		case <-time.After(time.Second * time.Duration(c.flagTimeout)):
			return fmt.Errorf("LXD still running after %ds timeout", c.flagTimeout)
		}
	}

	return <-chResult
}
<commit_msg>lxd/main/shutdown: Fix shutdown regression when running in snap<commit_after>package main

import (
	"fmt"
	"net/url"
	"strconv"
	"time"

	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
)

type cmdShutdown struct {
	global *cmdGlobal

	flagForce   bool
	flagTimeout int
}

func (c *cmdShutdown) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = "shutdown"
	cmd.Short = "Tell LXD to shutdown all containers and exit"
	cmd.Long = `Description:
  Tell LXD to shutdown all containers and exit

  This will tell LXD to start a clean shutdown of all containers,
  followed by having itself shutdown and exit.

  This can take quite a while as containers can take a long time to
  shutdown, especially if a non-standard timeout was configured for them.
`
	cmd.RunE = c.Run
	cmd.Flags().IntVarP(&c.flagTimeout, "timeout", "t", 0, "Number of seconds to wait before giving up"+"``")
	cmd.Flags().BoolVarP(&c.flagForce, "force", "f", false, "Force shutdown instead of waiting for running operations to finish"+"``")

	return cmd
}

func (c *cmdShutdown) Run(cmd *cobra.Command, args []string) error {
	connArgs := &lxd.ConnectionArgs{
		SkipGetServer: true,
	}

	d, err := lxd.ConnectLXDUnix("", connArgs)
	if err != nil {
		return err
	}

	v := url.Values{}
	v.Set("force", strconv.FormatBool(c.flagForce))

	chResult := make(chan error, 1)
	go func() {
		defer close(chResult)

		// Request shutdown, this shouldn't return until daemon has stopped.
		_, _, err = d.RawQuery("PUT", fmt.Sprintf("/internal/shutdown?%s", v.Encode()), nil, "")
		if err != nil {
			chResult <- err
			return
		}
	}()

	if c.flagTimeout > 0 {
		select {
		case err = <-chResult:
			return err
		case <-time.After(time.Second * time.Duration(c.flagTimeout)):
			return fmt.Errorf("LXD still running after %ds timeout", c.flagTimeout)
		}
	}

	return <-chResult
}
<|endoftext|>


<commit_before>package lzma

import (
	"errors"
	"fmt"
)

// decoderDict represents the dictionary while decoding LZMA files. It allows
// reading and decompression of the data at the same time.
type decoderDict struct {
	data []byte
	// history length
	h int
	// buffer length
	b int
	// current index
	c int
	// reader index
	r int
	// provides total length
	total int64
}

// newDecoderDict initializes a new decoderDict instance. If the arguments are
// negative or zero an error is returned.
func newDecoderDict(bufferLen int, historyLen int) (p *decoderDict, err error) {
	if !(0 < bufferLen) {
		return nil, errors.New("bufferLen must be positive")
	}
	if !(0 < historyLen) {
		return nil, errors.New("historyLen must be positive")
	}

	// k maximum match length
	k := maxLength
	if historyLen < k {
		k = historyLen
	}
	z := historyLen
	// We want to be able to copy the whole history, which is limited by
	// the reader index.
	if z <= maxLength {
		z += 1
	}
	// We want to ensure that there is enough place for a match to have
	// more than bufferLen bytes readable.
	if bufferLen+k > z {
		z = bufferLen + k
	}

	// check for overflows
	if z < bufferLen || z < historyLen {
		return nil, errors.New(
			"LZMA dictionary size overflows integer range")
	}

	p = &decoderDict{
		data: make([]byte, 0, z),
		h:    historyLen,
		b:    bufferLen,
	}
	return p, nil
}

// reset clears the decoder dictionary. This function must be called by the
// LZMA2 code if the decoder dictionary needs to be reset without a change in
// parameters. A change in parameters requires the dictionary to be newly
// initialized.
func (p *decoderDict) reset() {
	p.data = p.data[:0]
	p.c = 0
	p.r = 0
	p.total = 0
}

// readable returns the number of bytes available for reading. If it is bigger
// then buffer length then decompression should stop.
func (p *decoderDict) readable() int {
	delta := p.c - p.r
	if delta >= 0 {
		return delta
	}
	return len(p.data) + delta
}

// writable returns the number of bytes that can be currently written to the
// dictionary. The dictionary needs to be read to increase the number.
func (p *decoderDict) writable() int {
	delta := p.r - 1 - p.c
	if delta >= 0 {
		return delta
	}
	return cap(p.data) + delta
}

// Len returns the current length of the dictionary.
func (p *decoderDict) Len() int {
	if p.c > len(p.data) {
		panic("current index bigger then slice length")
	}
	if p.c == len(p.data) && p.c < p.h {
		return p.c
	}
	return p.h
}

// Read reads data from the dictionary into a. The function may return zero
// bytes.
func (p *decoderDict) Read(a []byte) (n int, err error) {
	if p.c < p.r {
		n = copy(a, p.data[p.r:])
		p.r += n
		if p.r < len(p.data) {
			return n, nil
		}
		p.r = 0
	}
	k := copy(a[n:], p.data[p.r:p.c])
	p.r += k
	n += k
	return n, nil
}

// errOverflow indicates an overflow situation of the dictionary. It can be
// addressed by reading more data from the dictionary.
var errOverflow = errors.New("overflow")

// Write puts the complete slice in the dictionary or no bytes at all. If no
// bytes are written an overflow will be indicated. The slice b is now allowed
// to overlap with the p.data slice to write to.
func (p *decoderDict) Write(b []byte) (n int, err error) {
	m := p.writable()
	n = len(b)
	if n > m {
		return 0, errOverflow
	}
	c := p.c + n
	z := cap(p.data)
	k := c
	if c > z {
		k = z
	}
	if len(p.data) < k {
		p.data = p.data[:k]
	}
	a := copy(p.data[p.c:], b)
	if a < n {
		p.c = copy(p.data, b[a:])
	} else {
		p.c = c
	}
	p.total += int64(n)
	return n, nil
}

// AddByte adds a single byte to the decoder dictionary. Even here it is
// possible that the write will overflow.
func (p *decoderDict) addByte(b byte) error {
	_, err := p.Write([]byte{b})
	return err
}

// CopyMatch copies a match with the given length n and distance d.
func (p *decoderDict) copyMatch(d, n int) error {
	if n <= 0 {
		return errors.New("argument n must be positive")
	}
	if d <= 0 {
		return errors.New("argument d must be positive")
	}
	if n > p.writable() {
		return errOverflow
	}
	if d > p.Len() {
		return errors.New("argument d is to large")
	}
	z := cap(p.data)
	i := p.c - d
	for n > 0 {
		a, b := i, i+n
		if b > p.c {
			b = p.c
		}
		if a < 0 {
			a, b = a+z, b+z
		}
		if b > z {
			b = z
		}
		k, err := p.Write(p.data[a:b])
		if err != nil {
			// Must panic here, because we don't want an incomplete
			// transaction
			panic(fmt.Sprintf("p.Write unexpected error %s", err))
		}
		n -= k
		i += k
	}
	return nil
}

// getByte returns the byte at distance d. If the distance is too large, the
// function returns zero.
func (p *decoderDict) getByte(d int) byte {
	if d < 0 {
		panic("negative d unexpected")
	}
	if d >= p.Len() {
		return 0
	}
	i := p.c - d
	if i < 0 {
		i += cap(p.data)
	}
	return p.data[i]
}
<commit_msg>lzma: error messages of copyMatch improved<commit_after>package lzma

import (
	"errors"
	"fmt"
)

// decoderDict represents the dictionary while decoding LZMA files. It allows
// reading and decompression of the data at the same time.
type decoderDict struct {
	data []byte
	// history length
	h int
	// buffer length
	b int
	// current index
	c int
	// reader index
	r int
	// provides total length
	total int64
}

// newDecoderDict initializes a new decoderDict instance. If the arguments are
// negative or zero an error is returned.
func newDecoderDict(bufferLen int, historyLen int) (p *decoderDict, err error) {
	if !(0 < bufferLen) {
		return nil, errors.New("bufferLen must be positive")
	}
	if !(0 < historyLen) {
		return nil, errors.New("historyLen must be positive")
	}

	// k maximum match length
	k := maxLength
	if historyLen < k {
		k = historyLen
	}
	z := historyLen
	// We want to be able to copy the whole history, which is limited by
	// the reader index.
	if z <= maxLength {
		z += 1
	}
	// We want to ensure that there is enough place for a match to have
	// more than bufferLen bytes readable.
	if bufferLen+k > z {
		z = bufferLen + k
	}

	// check for overflows
	if z < bufferLen || z < historyLen {
		return nil, errors.New(
			"LZMA dictionary size overflows integer range")
	}

	p = &decoderDict{
		data: make([]byte, 0, z),
		h:    historyLen,
		b:    bufferLen,
	}
	return p, nil
}

// reset clears the decoder dictionary. This function must be called by the
// LZMA2 code if the decoder dictionary needs to be reset without a change in
// parameters. A change in parameters requires the dictionary to be newly
// initialized.
func (p *decoderDict) reset() {
	p.data = p.data[:0]
	p.c = 0
	p.r = 0
	p.total = 0
}

// readable returns the number of bytes available for reading. If it is bigger
// then buffer length then decompression should stop.
func (p *decoderDict) readable() int {
	delta := p.c - p.r
	if delta >= 0 {
		return delta
	}
	return len(p.data) + delta
}

// writable returns the number of bytes that can be currently written to the
// dictionary. The dictionary needs to be read to increase the number.
func (p *decoderDict) writable() int {
	delta := p.r - 1 - p.c
	if delta >= 0 {
		return delta
	}
	return cap(p.data) + delta
}

// Len returns the current length of the dictionary.
func (p *decoderDict) Len() int {
	if p.c > len(p.data) {
		panic("current index bigger then slice length")
	}
	if p.c == len(p.data) && p.c < p.h {
		return p.c
	}
	return p.h
}

// Read reads data from the dictionary into a. The function may return zero
// bytes.
func (p *decoderDict) Read(a []byte) (n int, err error) {
	if p.c < p.r {
		n = copy(a, p.data[p.r:])
		p.r += n
		if p.r < len(p.data) {
			return n, nil
		}
		p.r = 0
	}
	k := copy(a[n:], p.data[p.r:p.c])
	p.r += k
	n += k
	return n, nil
}

// errOverflow indicates an overflow situation of the dictionary. It can be
// addressed by reading more data from the dictionary.
var errOverflow = errors.New("overflow")

// Write puts the complete slice in the dictionary or no bytes at all. If no
// bytes are written an overflow will be indicated. The slice b is now allowed
// to overlap with the p.data slice to write to.
func (p *decoderDict) Write(b []byte) (n int, err error) {
	m := p.writable()
	n = len(b)
	if n > m {
		return 0, errOverflow
	}
	c := p.c + n
	z := cap(p.data)
	k := c
	if c > z {
		k = z
	}
	if len(p.data) < k {
		p.data = p.data[:k]
	}
	a := copy(p.data[p.c:], b)
	if a < n {
		p.c = copy(p.data, b[a:])
	} else {
		p.c = c
	}
	p.total += int64(n)
	return n, nil
}

// AddByte adds a single byte to the decoder dictionary. Even here it is
// possible that the write will overflow.
func (p *decoderDict) addByte(b byte) error {
	_, err := p.Write([]byte{b})
	return err
}

// CopyMatch copies a match with the given length n and distance d.
func (p *decoderDict) copyMatch(d, n int) error {
	if n <= 0 {
		return errors.New("length n must be positive")
	}
	if d <= 0 {
		return errors.New("distance d must be positive")
	}
	if n > p.writable() {
		return errOverflow
	}
	if d > p.Len() {
		return errors.New("argument d is too large")
	}
	z := cap(p.data)
	i := p.c - d
	for n > 0 {
		a, b := i, i+n
		if b > p.c {
			b = p.c
		}
		if a < 0 {
			a, b = a+z, b+z
		}
		if b > z {
			b = z
		}
		k, err := p.Write(p.data[a:b])
		if err != nil {
			// Must panic here, because we don't want an incomplete
			// transaction
			panic(fmt.Sprintf("p.Write unexpected error %s", err))
		}
		n -= k
		i += k
	}
	return nil
}

// getByte returns the byte at distance d. If the distance is too large, the
// function returns zero.
func (p *decoderDict) getByte(d int) byte {
	if d < 0 {
		panic("negative d unexpected")
	}
	if d >= p.Len() {
		return 0
	}
	i := p.c - d
	if i < 0 {
		i += cap(p.data)
	}
	return p.data[i]
}
<|endoftext|>


<commit_before>//
// Copyright © 2014-2015 Exablox Corporation,  All Rights Reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
// COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
// ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
package main

import (
	"bufio"
	"fileutils"
	"flag"
	"fmt"
	"licensedb"
	"log"
	"notice"
	"os"
	"path/filepath"
	"runtime"
	"strutils"
	"sync"
	"tagger"
	"version"
)

type FileInfo struct {
	path string
	info os.FileInfo
}

type NoticeMsg struct {
	path   string
	notice *notice.Notice
}

const LicenseDBNumBuckets = 1000000

var ldb *licensedb.LicenseDB
var verbose bool
var ignoreErrors bool
var showLic bool
var quiet bool
var copyrightTagger *tagger.Tagger
var wg sync.WaitGroup
var workerChan chan FileInfo
var noticeChan chan NoticeMsg
var doneChan chan bool

const headHead = "<!DOCTYPE html>\n" +
	"\n" +
	"\n"

const headStyle = "	<meta charset=\"UTF-8\">\n" +
	"	\n"

const headFooter = "\n" +
	""

const footer = "\n" +
	"\n"

func ExtraUsage() {
	fmt.Printf("Usage: %s [options] [path] ...\n", os.Args[0])
	fmt.Printf("\n")
	fmt.Printf("Version %s. © 2014-2015 Exablox Corporation.  All Rights Reserved.\n", version.Version)
	fmt.Printf("\n")
	fmt.Printf("Options:\n")
	fmt.Printf("\n")
	flag.PrintDefaults()
	fmt.Printf("\n")
	fmt.Printf("Description:\n")
	fmt.Printf(`
  The purpose of this tool is to generate find and emit the licenses
  and copyright notices in a sourcetree for compliance purposes.
  For example, it can be used to help comply with the various open
  source licenses that require attribution in documentation.

  It recursively searches through the given set of files and
  directories, extracts copyright notices (de-duplicating against
  notices it has already seen), and extracts the licenses.

  At the end of the search, it outputs (-o) an HTML document which
  contains per file copyright notices and licenses.  If the -ldir
  option is given in combination with -o, it copies the licenses
  it finds into the specified directory, and makes the licenses
  viewable / downloadable via a link in the HTML document it emits.

  For directories provided on the command line, all contents will
  be recursively scanned.  For more complex searches (such as
  excluding certin file types), a list of files may also be provided
  via a file or stdin.  This makes it easy to build complex query
  pipelines with tools such as find(1).  See the '-i' and '-0'
  command line options for details.
`)
}

func handleParseError(path string, err error) error {
	if ignoreErrors {
		if !quiet {
			log.Printf("[ERROR] %s: %s\n", path, err)
		}
		return nil
	}

	return err
}

// Starts N workers determined by GOMAXPROCS
// And start a secretary worker
func setupWorkers() {
	var numWorkers int = runtime.GOMAXPROCS(-1)
	workerChan = make(chan FileInfo, numWorkers)
	noticeChan = make(chan NoticeMsg, numWorkers)
	doneChan = make(chan bool)

	// start up the one secretary that will record the notices
	go noticeHandler(noticeChan, doneChan)

	// start up the workers that make notices
	for i := 0; i < numWorkers; i++ {
		go fileHandler(workerChan, noticeChan, doneChan)
	}
}

// The worker that keeps trying to pull work from the worker channel then attempts
// to parse that info to create a notice then if a notice was created send that onto a
// channel to be logged in ldb
func fileHandler(workerChan chan FileInfo, noticeChan chan NoticeMsg, doneChan chan bool) {
	for {
		fileInfo, ok := <-workerChan // recieve the work to do and if work was recieved
		if !ok {                     // ok is set to false when workerChan is closed
			break
		}

		notice := FileParse(fileInfo.path, fileInfo.info)
		if notice != nil {
			noticeChan <- NoticeMsg{path: fileInfo.path, notice: notice}
		}
	}
	doneChan <- true
}

// A Secretart function that attempts to continually take notices and add them to the ldb
// once there are no more notices to take this stops
func noticeHandler(noticeChan chan NoticeMsg, doneChan chan bool) {
	for {
		noticeMsg, ok := <-noticeChan // recieve the notice made and if a notice was recieved
		if !ok {                      // ok is set to false when noticeChan is closed
			break
		}

		ldb.Add(noticeMsg.path, noticeMsg.notice, verbose)
	}
	doneChan <- true
}

// Gathers all the workers ensuring all work is done and all go routines have stopped before allowing
// the program to continue
func shutdownWorkers() {
	var numWorkers int = runtime.GOMAXPROCS(-1)
	for i := 0; i < numWorkers; i++ {
		<-doneChan
	}
	// I have revieved my N workers so no more notices will be comming in
	// close the notice chan so the Notice Handler can stop
	close(noticeChan)
	<-doneChan
	close(doneChan)
}

func sendWork(path string, info os.FileInfo, err error) error {
	if err != nil {
		return handleParseError(path, err)
	}

	workerChan <- FileInfo{path: path, info: info}

	return nil
}

//
// Directories:	filepath.Walk() handles, we'll be called again for each member
// Files:	Read and process
// Others:	Skip
//
func FileParse(path string, info os.FileInfo) *notice.Notice {
	if !info.Mode().IsRegular() {
		if verbose {
			log.Printf("[INFO] Skipping %s (not a file)\n", path)
		}
		return nil
	}

	lic, err := notice.NewNoticeFromFile(path, verbose, showLic, copyrightTagger)
	if err != nil {
		if ignoreErrors {
			if !quiet {
				log.Printf("[ERROR] %s: %s\n", path, err)
			}
		}
		log.Fatal(err)
	}

	return lic
}

// walks all files sending the path to sendWork
func ProcessFile(path string) error {

	err := filepath.Walk(path, sendWork)
	return err
}

//
// Workaround an OSX issue "regexec error 17, (illegal byte sequence)"
//
func fixenv() {
	v := os.Getenv("LANG")
	if v != "" && v != "C" {
		log.Printf("[WARN]: Overriding 'LANG=%s'\n", v)
	}

	err := os.Setenv("LANG", "C")
	if err != nil {
		log.Fatal(err)
	}

	if verbose {
		log.Printf("[WARN]: setenv 'LANG=C'\n")
	}
}

func main() {
	var inPath string
	var stylePath string
	var licenseDir string
	var corpusPath string
	var showVer bool
	var outPath string
	var zeroDelim bool

	flag.Usage = ExtraUsage

	flag.StringVar(&inPath, "i", "", "File to read list of files and directories from (use '-' for stdin)")
	flag.StringVar(&outPath, "o", "", "File to write HTML formatted licensedb to (default = stdout)")
	flag.StringVar(&licenseDir, "ldir", "", "Directory to save licenses to (default = don't save) ")

	flag.StringVar(&stylePath, "style", "", "Use this css stylesheet (default = embed)")
	flag.StringVar(&corpusPath, "corpus", "", "The path the the corpus to use for training the tagger model")

	flag.BoolVar(&zeroDelim, "0", false, "Pathnames read from the input file (-i) are \\0 delimited (default is \\n delimited)")
	flag.BoolVar(&ignoreErrors, "continue", false, "Continue processing, ignoring errors (default is abort on error)")
	flag.BoolVar(&quiet, "quiet", false, "Don't output errors (use in conjunction with '-continue')")
	flag.BoolVar(&verbose, "verbose", false, "Turn on verbose debug output (default is off)")
	flag.BoolVar(&showLic, "showlic", false, "show licenses found during processing")
	flag.BoolVar(&showVer, "version", false, "show version and exit")
	flag.Parse()

	if showVer {
		fmt.Printf("Version %s\n", version.Version)
		return
	}

	fixenv()

	var err error

	// Initialize the Tagger Model
	if corpusPath == "" {
		log.Fatal("corpus required")
		return
	}
	copyrightTagger = tagger.New(corpusPath)

	for _, path := range flag.Args() {
		err = fileutils.PathCheck(path, verbose)
		if err != nil {
			log.Fatal(err)
		}
	}

	ldb = licensedb.NewLicenseDB(licenseDir, LicenseDBNumBuckets, 0)

	setupWorkers()
	for _, path := range flag.Args() {
		err = ProcessFile(path)
		if err != nil {
			log.Fatal(err)
		}
	}
	close(workerChan) //done sending work
	shutdownWorkers()

	if inPath != "" {
		var infile *os.File
		var err error
		if inPath == "-" {
			infile = os.Stdin
		} else {
			infile, err = os.OpenFile(inPath, os.O_RDONLY, 0)
			if err != nil {
				log.Fatal(err)
			}
		}

		scanner := bufio.NewScanner(infile)

		if zeroDelim {
			scanner.Split(strutils.ScanZeros)
		}

		setupWorkers()
		for scanner.Scan() {
			err = ProcessFile(scanner.Text())
			if err != nil {
				log.Fatal(err)
			}
		}
		close(workerChan) //done sending work
		shutdownWorkers()

		err = scanner.Err()
		if err != nil {
			log.Fatal(err)
		}
		infile.Close()
	}

	outfile := os.Stdout
	if outPath != "" {
		outfile, err = os.OpenFile(outPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			log.Fatal(err)
		}
		defer outfile.Close()
	}

	outb := bufio.NewWriter(outfile)

	_, err = outb.WriteString(headHead)
	if err != nil {
		goto fail
	}
	if stylePath == "" {
		_, err = outb.WriteString(headStyle)
	} else {
		_, err = fmt.Fprintf(outb, "	<link rel=\"stylesheet\" type=\"text/css\" href=\"%s\">\n", stylePath)
	}
	if err != nil {
		goto fail
	}

	_, err = outb.WriteString(headFooter)
	if err != nil {
		goto fail
	}

	err = ldb.SortedSave(outb, verbose)

	if err != nil {
		goto fail
	}
	_, err = outb.WriteString(footer)
	if err != nil {
		goto fail
	}
	err = outb.Flush()
	if err != nil {
		goto fail
	}
	return

fail:
	outb.Flush()
	log.Fatal(err)
}
<commit_msg>Worker Setup/Takedown<commit_after>//
// Copyright © 2014-2015 Exablox Corporation,  All Rights Reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
// COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
// ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
package main

import (
	"bufio"
	"fileutils"
	"flag"
	"fmt"
	"licensedb"
	"log"
	"notice"
	"os"
	"path/filepath"
	"runtime"
	"strutils"
	"sync"
	"tagger"
	"version"
)

type FileInfo struct {
	path string
	info os.FileInfo
}

type NoticeMsg struct {
	path   string
	notice *notice.Notice
}

const LicenseDBNumBuckets = 1000000

var ldb *licensedb.LicenseDB
var verbose bool
var ignoreErrors bool
var showLic bool
var quiet bool
var copyrightTagger *tagger.Tagger
var wg sync.WaitGroup
var workerChan chan FileInfo
var noticeChan chan NoticeMsg
var doneChan chan bool

const headHead = "<!DOCTYPE html>\n" +
	"\n" +
	"\n"

const headStyle = "	<meta charset=\"UTF-8\">\n" +
	"	\n"

const headFooter = "\n" +
	""

const footer = "\n" +
	"\n"

func ExtraUsage() {
	fmt.Printf("Usage: %s [options] [path] ...\n", os.Args[0])
	fmt.Printf("\n")
	fmt.Printf("Version %s. © 2014-2015 Exablox Corporation.  All Rights Reserved.\n", version.Version)
	fmt.Printf("\n")
	fmt.Printf("Options:\n")
	fmt.Printf("\n")
	flag.PrintDefaults()
	fmt.Printf("\n")
	fmt.Printf("Description:\n")
	fmt.Printf(`
  The purpose of this tool is to generate find and emit the licenses
  and copyright notices in a sourcetree for compliance purposes.
  For example, it can be used to help comply with the various open
  source licenses that require attribution in documentation.

  It recursively searches through the given set of files and
  directories, extracts copyright notices (de-duplicating against
  notices it has already seen), and extracts the licenses.

  At the end of the search, it outputs (-o) an HTML document which
  contains per file copyright notices and licenses.  If the -ldir
  option is given in combination with -o, it copies the licenses
  it finds into the specified directory, and makes the licenses
  viewable / downloadable via a link in the HTML document it emits.

  For directories provided on the command line, all contents will
  be recursively scanned.  For more complex searches (such as
  excluding certin file types), a list of files may also be provided
  via a file or stdin.  This makes it easy to build complex query
  pipelines with tools such as find(1).  See the '-i' and '-0'
  command line options for details.
`)
}

func handleParseError(path string, err error) error {
	if ignoreErrors {
		if !quiet {
			log.Printf("[ERROR] %s: %s\n", path, err)
		}
		return nil
	}

	return err
}

// Starts N workers determined by GOMAXPROCS
// And start a secretary worker
func setupWorkers() {
	var numWorkers int = runtime.GOMAXPROCS(-1)
	workerChan = make(chan FileInfo, numWorkers)
	noticeChan = make(chan NoticeMsg, numWorkers)
	doneChan = make(chan bool)

	// start up the one secretary that will record the notices
	go noticeHandler(noticeChan, doneChan)

	// start up the workers that make notices
	for i := 0; i < numWorkers; i++ {
		go fileHandler(workerChan, noticeChan, doneChan)
	}
}

// The worker that keeps trying to pull work from the worker channel then attempts
// to parse that info to create a notice then if a notice was created send that onto a
// channel to be logged in ldb
func fileHandler(workerChan chan FileInfo, noticeChan chan NoticeMsg, doneChan chan bool) {
	for {
		fileInfo, ok := <-workerChan // recieve the work to do and if work was recieved
		if !ok {                     // ok is set to false when workerChan is closed
			break
		}

		notice := FileParse(fileInfo.path, fileInfo.info)
		if notice != nil {
			noticeChan <- NoticeMsg{path: fileInfo.path, notice: notice}
		}
	}
	doneChan <- true
}

// A Secretart function that attempts to continually take notices and add them to the ldb
// once there are no more notices to take this stops
func noticeHandler(noticeChan chan NoticeMsg, doneChan chan bool) {
	for {
		noticeMsg, ok := <-noticeChan // recieve the notice made and if a notice was recieved
		if !ok {                      // ok is set to false when noticeChan is closed
			break
		}

		ldb.Add(noticeMsg.path, noticeMsg.notice, verbose)
	}
	doneChan <- true
}

// Gathers all the workers ensuring all work is done and all go routines have stopped before allowing
// the program to continue
func shutdownWorkers() {
	close(workerChan) //done sending work
	var numWorkers int = runtime.GOMAXPROCS(-1)
	for i := 0; i < numWorkers; i++ {
		<-doneChan
	}
	// I have revieved my N workers so no more notices will be comming in
	// close the notice chan so the Notice Handler can stop
	close(noticeChan)
	<-doneChan
	close(doneChan)
}

func sendWork(path string, info os.FileInfo, err error) error {
	if err != nil {
		return handleParseError(path, err)
	}

	workerChan <- FileInfo{path: path, info: info}

	return nil
}

//
// Directories:	filepath.Walk() handles, we'll be called again for each member
// Files:	Read and process
// Others:	Skip
//
func FileParse(path string, info os.FileInfo) *notice.Notice {
	if !info.Mode().IsRegular() {
		if verbose {
			log.Printf("[INFO] Skipping %s (not a file)\n", path)
		}
		return nil
	}

	lic, err := notice.NewNoticeFromFile(path, verbose, showLic, copyrightTagger)
	if err != nil {
		if ignoreErrors {
			if !quiet {
				log.Printf("[ERROR] %s: %s\n", path, err)
			}
		}
		log.Fatal(err)
	}

	return lic
}

// walks all files sending the path to sendWork
func ProcessFile(path string) error {

	err := filepath.Walk(path, sendWork)
	return err
}

//
// Workaround an OSX issue "regexec error 17, (illegal byte sequence)"
//
func fixenv() {
	v := os.Getenv("LANG")
	if v != "" && v != "C" {
		log.Printf("[WARN]: Overriding 'LANG=%s'\n", v)
	}

	err := os.Setenv("LANG", "C")
	if err != nil {
		log.Fatal(err)
	}

	if verbose {
		log.Printf("[WARN]: setenv 'LANG=C'\n")
	}
}

func main() {
	var inPath string
	var stylePath string
	var licenseDir string
	var corpusPath string
	var showVer bool
	var outPath string
	var zeroDelim bool

	flag.Usage = ExtraUsage

	flag.StringVar(&inPath, "i", "", "File to read list of files and directories from (use '-' for stdin)")
	flag.StringVar(&outPath, "o", "", "File to write HTML formatted licensedb to (default = stdout)")
	flag.StringVar(&licenseDir, "ldir", "", "Directory to save licenses to (default = don't save) ")

	flag.StringVar(&stylePath, "style", "", "Use this css stylesheet (default = embed)")
	flag.StringVar(&corpusPath, "corpus", "", "The path the the corpus to use for training the tagger model")

	flag.BoolVar(&zeroDelim, "0", false, "Pathnames read from the input file (-i) are \\0 delimited (default is \\n delimited)")
	flag.BoolVar(&ignoreErrors, "continue", false, "Continue processing, ignoring errors (default is abort on error)")
	flag.BoolVar(&quiet, "quiet", false, "Don't output errors (use in conjunction with '-continue')")
	flag.BoolVar(&verbose, "verbose", false, "Turn on verbose debug output (default is off)")
	flag.BoolVar(&showLic, "showlic", false, "show licenses found during processing")
	flag.BoolVar(&showVer, "version", false, "show version and exit")
	flag.Parse()

	if showVer {
		fmt.Printf("Version %s\n", version.Version)
		return
	}

	fixenv()

	var err error

	// Initialize the Tagger Model
	if corpusPath == "" {
		log.Fatal("corpus required")
		return
	}
	copyrightTagger = tagger.New(corpusPath)

	for _, path := range flag.Args() {
		err = fileutils.PathCheck(path, verbose)
		if err != nil {
			log.Fatal(err)
		}
	}

	ldb = licensedb.NewLicenseDB(licenseDir, LicenseDBNumBuckets, 0)

	setupWorkers()
	for _, path := range flag.Args() {
		err = ProcessFile(path)
		if err != nil {
			log.Fatal(err)
		}
	}

	if inPath != "" {
		var infile *os.File
		var err error
		if inPath == "-" {
			infile = os.Stdin
		} else {
			infile, err = os.OpenFile(inPath, os.O_RDONLY, 0)
			if err != nil {
				log.Fatal(err)
			}
		}

		scanner := bufio.NewScanner(infile)

		if zeroDelim {
			scanner.Split(strutils.ScanZeros)
		}

		for scanner.Scan() {
			err = ProcessFile(scanner.Text())
			if err != nil {
				log.Fatal(err)
			}
		}

		err = scanner.Err()
		if err != nil {
			log.Fatal(err)
		}
		infile.Close()
	}
	shutdownWorkers()

	outfile := os.Stdout
	if outPath != "" {
		outfile, err = os.OpenFile(outPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			log.Fatal(err)
		}
		defer outfile.Close()
	}

	outb := bufio.NewWriter(outfile)

	_, err = outb.WriteString(headHead)
	if err != nil {
		goto fail
	}
	if stylePath == "" {
		_, err = outb.WriteString(headStyle)
	} else {
		_, err = fmt.Fprintf(outb, "	<link rel=\"stylesheet\" type=\"text/css\" href=\"%s\">\n", stylePath)
	}
	if err != nil {
		goto fail
	}

	_, err = outb.WriteString(headFooter)
	if err != nil {
		goto fail
	}

	err = ldb.SortedSave(outb, verbose)

	if err != nil {
		goto fail
	}
	_, err = outb.WriteString(footer)
	if err != nil {
		goto fail
	}
	err = outb.Flush()
	if err != nil {
		goto fail
	}
	return

fail:
	outb.Flush()
	log.Fatal(err)
}
<|endoftext|>


<commit_before>// A quick and easy way to setup a RESTful JSON API
//
// Go-Json-Rest is a thin layer on top of net/http that helps building RESTful JSON APIs easily.
// It provides fast URL routing using https://github.com/ant0ine/go-urlrouter, and helpers to deal
// with JSON requests and responses. It is not a high-level REST framework that transparently maps
// HTTP requests to procedure calls, on the opposite, you constantly have access to the underlying
// net/http objects.
//
// Example:
//
//      package main
//
//      import (
//              "github.com/ant0ine/go-json-rest"
//              "net/http"
//      )
//
//      type User struct {
//              Id   string
//              Name string
//      }
//
//      func GetUser(w *rest.ResponseWriter, req *rest.Request) {
//              user := User{
//                      Id:   req.PathParam("id"),
//                      Name: "Antoine",
//              }
//              w.WriteJson(&user)
//      }
//
//      func main() {
//              handler := ResourceHandler{}
//              handler.SetRoutes(
//                      rest.Route{"GET", "/users/:id", GetUser},
//              )
//              http.ListenAndServe(":8080", &handler)
//      }
//
package rest

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"github.com/ant0ine/go-urlrouter"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"reflect"
	"runtime/debug"
	"strings"
	"time"
)

// Implement the http.Handler interface and act as a router for the defined Routes.
// The defaults are intended to be developemnt friendly, for production you may want
// to turn on gzip and disable the JSON indentation.
type ResourceHandler struct {
	router urlrouter.Router

	// If true, and if the client accepts the Gzip encoding, the response payloads
	// will be compressed using gzip, and the corresponding response header will set.
	EnableGzip bool

	// If true, the JSON payload will be written in one line with no space.
	DisableJsonIndent bool

	// If true, when a "panic" happens, the error string and the stack trace will be
	// printed in the 500 response body.
	EnableResponseStackTrace bool

	// If true, the record that is logged for each response will be printed as JSON
	// in the log. Convenient for log parsing.
	EnableLogAsJson bool

	// Custom logger, defaults to log.New(os.Stderr, "", log.LstdFlags)
	Logger *log.Logger
}

// Used with SetRoutes.
type Route struct {
	HttpMethod string
	PathExp    string
	Func       func(*ResponseWriter, *Request)
}

// Create a Route that points to an object method. It can be convenient to point to an object method instead
// of a function, this helper makes it easy by passing the object instance and the method name as parameters.
func RouteObjectMethod(http_method string, path_exp string, object_instance interface{}, object_method string) Route {

	value := reflect.ValueOf(object_instance)
	func_value := value.MethodByName(object_method)
	if func_value.IsValid() == false {
		panic(fmt.Sprintf(
			"Cannot find the object method %s on %s",
			object_method,
			value,
		))
	}
	route_func := func(w *ResponseWriter, r *Request) {
		func_value.Call([]reflect.Value{
			reflect.ValueOf(w),
			reflect.ValueOf(r),
		})
	}

	return Route{
		HttpMethod: http_method,
		PathExp:    path_exp,
		Func:       route_func,
	}
}

// Define the Routes. The order the Routes matters,
// if a request matches multiple Routes, the first one will be used.
// Note that the underlying router is https://github.com/ant0ine/go-urlrouter.
func (self *ResourceHandler) SetRoutes(routes ...Route) error {
	self.router = urlrouter.Router{
		Routes: []urlrouter.Route{},
	}
	for _, route := range routes {
		// make sure the method is uppercase
		http_method := strings.ToUpper(route.HttpMethod)

		self.router.Routes = append(
			self.router.Routes,
			urlrouter.Route{
				PathExp: http_method + route.PathExp,
				Dest:    route.Func,
			},
		)
	}
	return self.router.Start()
}

type response_log_record struct {
	StatusCode   int
	ResponseTime *time.Duration
	HttpMethod   string
	RequestURI   string
}

func (self *ResourceHandler) log_response(record *response_log_record) {
	if self.EnableLogAsJson {
		b, err := json.Marshal(record)
		if err != nil {
			panic(err)
		}
		self.Logger.Printf("%s", b)
	} else {
		self.Logger.Printf("%d %v %s %s",
			record.StatusCode,
			record.ResponseTime,
			record.HttpMethod,
			record.RequestURI,
		)
	}
}

// This makes ResourceHandler implement the http.Handler interface.
// You probably don't want to use it directly.
func (self *ResourceHandler) ServeHTTP(orig_writer http.ResponseWriter, orig_request *http.Request) {

	t0 := time.Now()

	// set a default Logger
	if self.Logger == nil {
		self.Logger = log.New(os.Stderr, "", log.LstdFlags)
	}

	// catch user code's panic, and convert to http response
	defer func() {
		if r := recover(); r != nil {
			trace := debug.Stack()

			// log the trace
			self.Logger.Printf("%s\n%s", r, trace)

			// write error response
			message := "Internal Server Error"
			if self.EnableResponseStackTrace {
				message = fmt.Sprintf("%s\n\n%s", r, trace)
			}
			http.Error(orig_writer, message, http.StatusInternalServerError)

			// log response
			self.log_response(&response_log_record{
				http.StatusNotFound,
				nil,
				orig_request.Method,
				orig_request.URL.RequestURI(),
			})
		}
	}()

	// find the route
	route, params, err := self.router.FindRoute(
		strings.ToUpper(orig_request.Method) + orig_request.URL.Path,
	)
	if err != nil {
		// should never happen as the URL has already been parsed
		panic(err)
	}
	if route == nil {
		// no route found
		http.NotFound(orig_writer, orig_request)

		// log response
		self.log_response(&response_log_record{
			http.StatusNotFound,
			nil,
			orig_request.Method,
			orig_request.URL.RequestURI(),
		})
		return
	}

	// determine if gzip is needed
	is_gzipped := self.EnableGzip == true &&
		strings.Contains(orig_request.Header.Get("Accept-Encoding"), "gzip")

	is_indented := !self.DisableJsonIndent

	request := Request{
		orig_request,
		params,
	}

	writer := ResponseWriter{
		orig_writer,
		is_gzipped,
		is_indented,
		0,
		false,
	}

	// run the user code
	handler := route.Dest.(func(*ResponseWriter, *Request))
	handler(&writer, &request)

	duration := time.Now().Sub(t0)

	// log response
	self.log_response(&response_log_record{
		writer.status_code,
		&duration,
		orig_request.Method,
		orig_request.URL.RequestURI(),
	})
}

// Inherit from http.Request, and provide additional methods.
type Request struct {
	*http.Request
	// map of parameters that have been matched in the URL Path.
	PathParams map[string]string
}

// Provide a convenient access to the PathParams map
func (self *Request) PathParam(name string) string {
	return self.PathParams[name]
}

// Read the request body and decode the JSON using json.Unmarshal
func (self *Request) DecodeJsonPayload(v interface{}) error {
	content, err := ioutil.ReadAll(self.Body)
	self.Body.Close()
	if err != nil {
		return err
	}
	err = json.Unmarshal(content, v)
	if err != nil {
		return err
	}
	return nil
}

// Inherit from an object implementing the http.ResponseWriter interface, and provide additional methods.
type ResponseWriter struct {
	http.ResponseWriter
	is_gzipped   bool
	is_indented  bool
	status_code  int
	wrote_header bool
}

// Overloading of the http.ResponseWriter method.
// Just record the status code for logging.
func (self *ResponseWriter) WriteHeader(code int) {
	self.ResponseWriter.WriteHeader(code)
	self.status_code = code
	self.wrote_header = true
}

// Overloading of the http.ResponseWriter method.
// Provide additional capabilities, like transparent gzip encoding.
func (self *ResponseWriter) Write(b []byte) (int, error) {

	if !self.wrote_header {
		self.WriteHeader(http.StatusOK)
	}

	if self.is_gzipped {
		self.Header().Set("Content-Encoding", "gzip")
		gzip_writer := gzip.NewWriter(self.ResponseWriter)
		defer gzip_writer.Close()
		return gzip_writer.Write(b)
	}

	return self.ResponseWriter.Write(b)
}

// Encode the object in JSON, set the content-type header,
// and call Write
func (self *ResponseWriter) WriteJson(v interface{}) error {
	self.Header().Set("content-type", "application/json")
	var b []byte
	var err error
	if self.is_indented {
		b, err = json.MarshalIndent(v, "", "  ")
	} else {
		b, err = json.Marshal(v)
	}
	if err != nil {
		return err
	}
	self.Write(b)
	return nil
}
<commit_msg>make sure to log the response time in all cases<commit_after>// A quick and easy way to setup a RESTful JSON API
//
// Go-Json-Rest is a thin layer on top of net/http that helps building RESTful JSON APIs easily.
// It provides fast URL routing using https://github.com/ant0ine/go-urlrouter, and helpers to deal
// with JSON requests and responses. It is not a high-level REST framework that transparently maps
// HTTP requests to procedure calls, on the opposite, you constantly have access to the underlying
// net/http objects.
//
// Example:
//
//      package main
//
//      import (
//              "github.com/ant0ine/go-json-rest"
//              "net/http"
//      )
//
//      type User struct {
//              Id   string
//              Name string
//      }
//
//      func GetUser(w *rest.ResponseWriter, req *rest.Request) {
//              user := User{
//                      Id:   req.PathParam("id"),
//                      Name: "Antoine",
//              }
//              w.WriteJson(&user)
//      }
//
//      func main() {
//              handler := ResourceHandler{}
//              handler.SetRoutes(
//                      rest.Route{"GET", "/users/:id", GetUser},
//              )
//              http.ListenAndServe(":8080", &handler)
//      }
//
package rest

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"github.com/ant0ine/go-urlrouter"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"reflect"
	"runtime/debug"
	"strings"
	"time"
)

// Implement the http.Handler interface and act as a router for the defined Routes.
// The defaults are intended to be developemnt friendly, for production you may want
// to turn on gzip and disable the JSON indentation.
type ResourceHandler struct {
	router urlrouter.Router

	// If true, and if the client accepts the Gzip encoding, the response payloads
	// will be compressed using gzip, and the corresponding response header will set.
	EnableGzip bool

	// If true, the JSON payload will be written in one line with no space.
	DisableJsonIndent bool

	// If true, when a "panic" happens, the error string and the stack trace will be
	// printed in the 500 response body.
	EnableResponseStackTrace bool

	// If true, the record that is logged for each response will be printed as JSON
	// in the log. Convenient for log parsing.
	EnableLogAsJson bool

	// Custom logger, defaults to log.New(os.Stderr, "", log.LstdFlags)
	Logger *log.Logger
}

// Used with SetRoutes.
type Route struct {
	HttpMethod string
	PathExp    string
	Func       func(*ResponseWriter, *Request)
}

// Create a Route that points to an object method. It can be convenient to point to an object method instead
// of a function, this helper makes it easy by passing the object instance and the method name as parameters.
func RouteObjectMethod(http_method string, path_exp string, object_instance interface{}, object_method string) Route {

	value := reflect.ValueOf(object_instance)
	func_value := value.MethodByName(object_method)
	if func_value.IsValid() == false {
		panic(fmt.Sprintf(
			"Cannot find the object method %s on %s",
			object_method,
			value,
		))
	}
	route_func := func(w *ResponseWriter, r *Request) {
		func_value.Call([]reflect.Value{
			reflect.ValueOf(w),
			reflect.ValueOf(r),
		})
	}

	return Route{
		HttpMethod: http_method,
		PathExp:    path_exp,
		Func:       route_func,
	}
}

// Define the Routes. The order the Routes matters,
// if a request matches multiple Routes, the first one will be used.
// Note that the underlying router is https://github.com/ant0ine/go-urlrouter.
func (self *ResourceHandler) SetRoutes(routes ...Route) error {
	self.router = urlrouter.Router{
		Routes: []urlrouter.Route{},
	}
	for _, route := range routes {
		// make sure the method is uppercase
		http_method := strings.ToUpper(route.HttpMethod)

		self.router.Routes = append(
			self.router.Routes,
			urlrouter.Route{
				PathExp: http_method + route.PathExp,
				Dest:    route.Func,
			},
		)
	}
	return self.router.Start()
}

type response_log_record struct {
	StatusCode   int
	ResponseTime *time.Duration
	HttpMethod   string
	RequestURI   string
}

func (self *ResourceHandler) log_response_record(record *response_log_record) {
	if self.EnableLogAsJson {
		b, err := json.Marshal(record)
		if err != nil {
			panic(err)
		}
		self.Logger.Printf("%s", b)
	} else {
		self.Logger.Printf("%d %v %s %s",
			record.StatusCode,
			record.ResponseTime,
			record.HttpMethod,
			record.RequestURI,
		)
	}
}

func (self *ResourceHandler) log_response(status_code int, start *time.Time, request *http.Request) {

	duration := time.Now().Sub(*start)

	self.log_response_record(&response_log_record{
		status_code,
		&duration,
		request.Method,
		request.URL.RequestURI(),
	})
}

// This makes ResourceHandler implement the http.Handler interface.
// You probably don't want to use it directly.
func (self *ResourceHandler) ServeHTTP(orig_writer http.ResponseWriter, orig_request *http.Request) {

	start := time.Now()

	// set a default Logger
	if self.Logger == nil {
		self.Logger = log.New(os.Stderr, "", log.LstdFlags)
	}

	// catch user code's panic, and convert to http response
	defer func() {
		if r := recover(); r != nil {
			trace := debug.Stack()

			// log the trace
			self.Logger.Printf("%s\n%s", r, trace)

			// write error response
			message := "Internal Server Error"
			if self.EnableResponseStackTrace {
				message = fmt.Sprintf("%s\n\n%s", r, trace)
			}
			http.Error(orig_writer, message, http.StatusInternalServerError)

			// log response
			self.log_response(
				http.StatusNotFound,
				&start,
				orig_request,
			)
		}
	}()

	// find the route
	route, params, err := self.router.FindRoute(
		strings.ToUpper(orig_request.Method) + orig_request.URL.Path,
	)
	if err != nil {
		// should never happen as the URL has already been parsed
		panic(err)
	}
	if route == nil {
		// no route found
		http.NotFound(orig_writer, orig_request)

		// log response
		self.log_response(
			http.StatusNotFound,
			&start,
			orig_request,
		)
		return
	}

	// determine if gzip is needed
	is_gzipped := self.EnableGzip == true &&
		strings.Contains(orig_request.Header.Get("Accept-Encoding"), "gzip")

	is_indented := !self.DisableJsonIndent

	request := Request{
		orig_request,
		params,
	}

	writer := ResponseWriter{
		orig_writer,
		is_gzipped,
		is_indented,
		0,
		false,
	}

	// run the user code
	handler := route.Dest.(func(*ResponseWriter, *Request))
	handler(&writer, &request)

	// log response
	self.log_response(
		writer.status_code,
		&start,
		orig_request,
	)
}

// Inherit from http.Request, and provide additional methods.
type Request struct {
	*http.Request
	// map of parameters that have been matched in the URL Path.
	PathParams map[string]string
}

// Provide a convenient access to the PathParams map
func (self *Request) PathParam(name string) string {
	return self.PathParams[name]
}

// Read the request body and decode the JSON using json.Unmarshal
func (self *Request) DecodeJsonPayload(v interface{}) error {
	content, err := ioutil.ReadAll(self.Body)
	self.Body.Close()
	if err != nil {
		return err
	}
	err = json.Unmarshal(content, v)
	if err != nil {
		return err
	}
	return nil
}

// Inherit from an object implementing the http.ResponseWriter interface, and provide additional methods.
type ResponseWriter struct {
	http.ResponseWriter
	is_gzipped   bool
	is_indented  bool
	status_code  int
	wrote_header bool
}

// Overloading of the http.ResponseWriter method.
// Just record the status code for logging.
func (self *ResponseWriter) WriteHeader(code int) {
	self.ResponseWriter.WriteHeader(code)
	self.status_code = code
	self.wrote_header = true
}

// Overloading of the http.ResponseWriter method.
// Provide additional capabilities, like transparent gzip encoding.
func (self *ResponseWriter) Write(b []byte) (int, error) {

	if !self.wrote_header {
		self.WriteHeader(http.StatusOK)
	}

	if self.is_gzipped {
		self.Header().Set("Content-Encoding", "gzip")
		gzip_writer := gzip.NewWriter(self.ResponseWriter)
		defer gzip_writer.Close()
		return gzip_writer.Write(b)
	}

	return self.ResponseWriter.Write(b)
}

// Encode the object in JSON, set the content-type header,
// and call Write
func (self *ResponseWriter) WriteJson(v interface{}) error {
	self.Header().Set("content-type", "application/json")
	var b []byte
	var err error
	if self.is_indented {
		b, err = json.MarshalIndent(v, "", "  ")
	} else {
		b, err = json.Marshal(v)
	}
	if err != nil {
		return err
	}
	self.Write(b)
	return nil
}
<|endoftext|>


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

package main

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

	"github.com/blevesearch/bleve"
	bleveHttp "github.com/blevesearch/bleve/http"
	"github.com/gorilla/mux"

	log "github.com/couchbaselabs/clog"
)

func NewManagerRESTRouter(mgr *Manager, staticDir string, mr *MsgRing) (*mux.Router, error) {
	// create a router to serve static files
	r := staticFileRouter(staticDir, []string{
		"/overview",
		"/search",
		"/indexes",
		"/analysis",
		"/monitor",
		"/logs",
	})

	r.Handle("/api/log", NewGetLogHandler(mr))
	r.Handle("/api/index", NewListIndexHandler(mgr)).Methods("GET")
	r.Handle("/api/index/{indexName}", NewCreateIndexHandler(mgr)).Methods("PUT")
	r.Handle("/api/index/{indexName}", NewDeleteIndexHandler(mgr)).Methods("DELETE")
	r.Handle("/api/index/{indexName}", NewGetIndexHandler(mgr)).Methods("GET")

	tags := mgr.Tags()

	if tags == nil || tags["queryer"] {
		r.Handle("/api/index/{indexName}/count", NewCountHandler(mgr)).Methods("GET")
		r.Handle("/api/index/{indexName}/search", NewSearchHandler(mgr)).Methods("POST")
	}

	// the rest are standard bleveHttp handlers for the lower "pindex" level...
	if tags == nil || tags["pindex"] {
		r.Handle("/api/pindex", bleveHttp.NewListIndexesHandler()).Methods("GET")

		getIndexHandler := bleveHttp.NewGetIndexHandler()
		getIndexHandler.IndexNameLookup = indexNameLookup
		r.Handle("/api/pindex/{indexName}", getIndexHandler).Methods("GET")

		docCountHandler := bleveHttp.NewDocCountHandler("")
		docCountHandler.IndexNameLookup = indexNameLookup
		r.Handle("/api/pindex/{indexName}/count", docCountHandler).Methods("GET")

		docGetHandler := bleveHttp.NewDocGetHandler("")
		docGetHandler.IndexNameLookup = indexNameLookup
		docGetHandler.DocIDLookup = docIDLookup
		r.Handle("/api/pindex/{indexName}/doc/{docID}", docGetHandler).Methods("GET")

		debugDocHandler := bleveHttp.NewDebugDocumentHandler("")
		debugDocHandler.IndexNameLookup = indexNameLookup
		debugDocHandler.DocIDLookup = docIDLookup
		r.Handle("/api/pindex/{indexName}/docDebug/{docID}", debugDocHandler).Methods("GET")

		searchHandler := bleveHttp.NewSearchHandler("")
		searchHandler.IndexNameLookup = indexNameLookup
		r.Handle("/api/pindex/{indexName}/search", searchHandler).Methods("POST")

		listFieldsHandler := bleveHttp.NewListFieldsHandler("")
		listFieldsHandler.IndexNameLookup = indexNameLookup
		r.Handle("/api/pindex/{indexName}/fields", listFieldsHandler).Methods("GET")

		r.Handle("/api/feedStats", NewFeedStatsHandler(mgr)).Methods("GET")
	}

	return r, nil
}

func muxVariableLookup(req *http.Request, name string) string {
	return mux.Vars(req)[name]
}

func docIDLookup(req *http.Request) string {
	return muxVariableLookup(req, "docID")
}

func indexNameLookup(req *http.Request) string {
	return muxVariableLookup(req, "indexName")
}

func indexAlias(mgr *Manager, indexName, indexUUID string) (bleve.IndexAlias, error) {
	// TODO: also add remote pindexes to alias, not just local pindexes.
	alias := bleve.NewIndexAlias()

	_, pindexes := mgr.CurrentMaps()
	for _, pindex := range pindexes {
		if pindex.IndexType == "bleve" &&
			pindex.IndexName == indexName &&
			(indexUUID == "" || pindex.IndexUUID == indexUUID) {
			bindex, ok := pindex.Impl.(bleve.Index)
			if ok && bindex != nil {
				alias.Add(bindex)
			}
		}
	}

	return alias, nil
}

// ---------------------------------------------------

type ListIndexHandler struct {
	mgr *Manager
}

func NewListIndexHandler(mgr *Manager) *ListIndexHandler {
	return &ListIndexHandler{mgr: mgr}
}

func (h *ListIndexHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexDefs, _, err := h.mgr.GetIndexDefs(false)
	if err != nil {
		showError(w, req, "could not retrieve index defs", 500)
		return
	}

	rv := struct {
		Status    string     `json:"status"`
		IndexDefs *IndexDefs `json:"indexDefs"`
	}{
		Status:    "ok",
		IndexDefs: indexDefs,
	}
	mustEncode(w, rv)
}

// ---------------------------------------------------

type GetIndexHandler struct {
	mgr *Manager
}

func NewGetIndexHandler(mgr *Manager) *GetIndexHandler {
	return &GetIndexHandler{mgr: mgr}
}

func (h *GetIndexHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexName := indexNameLookup(req)
	if indexName == "" {
		showError(w, req, "index name is required", 400)
		return
	}

	_, indexDefsByName, err := h.mgr.GetIndexDefs(false)
	if err != nil {
		showError(w, req, "could not retrieve index defs", 500)
		return
	}

	indexDef, exists := indexDefsByName[indexName]
	if !exists || indexDef == nil {
		showError(w, req, "not an index", 400)
		return
	}

	indexUUID := req.FormValue("indexUUID")
	if indexUUID != "" && indexUUID != indexDef.UUID {
		showError(w, req, "wrong index UUID", 400)
		return
	}

	m := map[string]interface{}{}
	if indexDef.Schema != "" {
		if err := json.Unmarshal([]byte(indexDef.Schema), &m); err != nil {
			showError(w, req, "could not unmarshal mapping", 500)
			return
		}
	}

	rv := struct {
		Status       string                 `json:"status"`
		IndexDef     *IndexDef              `json:"indexDef"`
		IndexMapping map[string]interface{} `json:"indexMapping"`
	}{
		Status:       "ok",
		IndexDef:     indexDef,
		IndexMapping: m,
	}
	mustEncode(w, rv)
}

// ---------------------------------------------------

type CountHandler struct {
	mgr *Manager
}

func NewCountHandler(mgr *Manager) *CountHandler {
	return &CountHandler{mgr: mgr}
}

func (h *CountHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexName := indexNameLookup(req)
	if indexName == "" {
		showError(w, req, "index name is required", 400)
		return
	}

	indexUUID := req.FormValue("indexUUID")

	alias, err := indexAlias(h.mgr, indexName, indexUUID)
	if err != nil {
		showError(w, req, fmt.Sprintf("index alias: %v", err), 500)
		return
	}

	docCount, err := alias.DocCount()
	if err != nil {
		showError(w, req, fmt.Sprintf("error counting docs: %v", err), 500)
		return
	}

	rv := struct {
		Status string `json:"status"`
		Count  uint64 `json:"count"`
	}{
		Status: "ok",
		Count:  docCount,
	}
	mustEncode(w, rv)
}

// ---------------------------------------------------

type SearchHandler struct {
	mgr *Manager
}

func NewSearchHandler(mgr *Manager) *SearchHandler {
	return &SearchHandler{mgr: mgr}
}

func (h *SearchHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexName := indexNameLookup(req)
	if indexName == "" {
		showError(w, req, "index name is required", 400)
		return
	}

	indexUUID := req.FormValue("indexUUID")

	log.Printf("search request: %s", indexName)

	alias, err := indexAlias(h.mgr, indexName, indexUUID)
	if err != nil {
		showError(w, req, fmt.Sprintf("index alias: %v", err), 500)
		return
	}

	// read the request body
	requestBody, err := ioutil.ReadAll(req.Body)
	if err != nil {
		showError(w, req, fmt.Sprintf("error reading request body: %v", err), 400)
		return
	}

	log.Printf("search request body: %s", requestBody)

	// parse the request
	var searchRequest bleve.SearchRequest
	err = json.Unmarshal(requestBody, &searchRequest)
	if err != nil {
		showError(w, req, fmt.Sprintf("error parsing query: %v", err), 400)
		return
	}

	log.Printf("search parsed request %#v", searchRequest)

	// varlidate the query
	err = searchRequest.Query.Validate()
	if err != nil {
		showError(w, req, fmt.Sprintf("error validating query: %v", err), 400)
		return
	}

	// execute the query
	searchResponse, err := alias.Search(&searchRequest)
	if err != nil {
		showError(w, req, fmt.Sprintf("error executing query: %v", err), 500)
		return
	}

	// encode the response
	mustEncode(w, searchResponse)
}

// ---------------------------------------------------

type FeedStatsHandler struct {
	mgr *Manager
}

func NewFeedStatsHandler(mgr *Manager) *FeedStatsHandler {
	return &FeedStatsHandler{mgr: mgr}
}

func (h *FeedStatsHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	feeds, _ := h.mgr.CurrentMaps()
	w.Write([]byte("[\n"))
	first := true
	for feedName, feed := range feeds {
		if !first {
			w.Write([]byte(",\n"))
		}
		first = false
		w.Write([]byte(fmt.Sprintf("  {\"feedName\":\"%s\",\"stats\":", feedName)))
		feed.Stats(w)
		w.Write([]byte("}\n"))
	}
	w.Write([]byte("]\n"))
}
<commit_msg>sorting feed stats by feedName<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"sort"

	"github.com/blevesearch/bleve"
	bleveHttp "github.com/blevesearch/bleve/http"
	"github.com/gorilla/mux"

	log "github.com/couchbaselabs/clog"
)

func NewManagerRESTRouter(mgr *Manager, staticDir string, mr *MsgRing) (*mux.Router, error) {
	// create a router to serve static files
	r := staticFileRouter(staticDir, []string{
		"/overview",
		"/search",
		"/indexes",
		"/analysis",
		"/monitor",
		"/logs",
	})

	r.Handle("/api/log", NewGetLogHandler(mr))
	r.Handle("/api/index", NewListIndexHandler(mgr)).Methods("GET")
	r.Handle("/api/index/{indexName}", NewCreateIndexHandler(mgr)).Methods("PUT")
	r.Handle("/api/index/{indexName}", NewDeleteIndexHandler(mgr)).Methods("DELETE")
	r.Handle("/api/index/{indexName}", NewGetIndexHandler(mgr)).Methods("GET")

	tags := mgr.Tags()

	if tags == nil || tags["queryer"] {
		r.Handle("/api/index/{indexName}/count", NewCountHandler(mgr)).Methods("GET")
		r.Handle("/api/index/{indexName}/search", NewSearchHandler(mgr)).Methods("POST")
	}

	// the rest are standard bleveHttp handlers for the lower "pindex" level...
	if tags == nil || tags["pindex"] {
		r.Handle("/api/pindex", bleveHttp.NewListIndexesHandler()).Methods("GET")

		getIndexHandler := bleveHttp.NewGetIndexHandler()
		getIndexHandler.IndexNameLookup = indexNameLookup
		r.Handle("/api/pindex/{indexName}", getIndexHandler).Methods("GET")

		docCountHandler := bleveHttp.NewDocCountHandler("")
		docCountHandler.IndexNameLookup = indexNameLookup
		r.Handle("/api/pindex/{indexName}/count", docCountHandler).Methods("GET")

		docGetHandler := bleveHttp.NewDocGetHandler("")
		docGetHandler.IndexNameLookup = indexNameLookup
		docGetHandler.DocIDLookup = docIDLookup
		r.Handle("/api/pindex/{indexName}/doc/{docID}", docGetHandler).Methods("GET")

		debugDocHandler := bleveHttp.NewDebugDocumentHandler("")
		debugDocHandler.IndexNameLookup = indexNameLookup
		debugDocHandler.DocIDLookup = docIDLookup
		r.Handle("/api/pindex/{indexName}/docDebug/{docID}", debugDocHandler).Methods("GET")

		searchHandler := bleveHttp.NewSearchHandler("")
		searchHandler.IndexNameLookup = indexNameLookup
		r.Handle("/api/pindex/{indexName}/search", searchHandler).Methods("POST")

		listFieldsHandler := bleveHttp.NewListFieldsHandler("")
		listFieldsHandler.IndexNameLookup = indexNameLookup
		r.Handle("/api/pindex/{indexName}/fields", listFieldsHandler).Methods("GET")

		r.Handle("/api/feedStats", NewFeedStatsHandler(mgr)).Methods("GET")
	}

	return r, nil
}

func muxVariableLookup(req *http.Request, name string) string {
	return mux.Vars(req)[name]
}

func docIDLookup(req *http.Request) string {
	return muxVariableLookup(req, "docID")
}

func indexNameLookup(req *http.Request) string {
	return muxVariableLookup(req, "indexName")
}

func indexAlias(mgr *Manager, indexName, indexUUID string) (bleve.IndexAlias, error) {
	// TODO: also add remote pindexes to alias, not just local pindexes.
	alias := bleve.NewIndexAlias()

	_, pindexes := mgr.CurrentMaps()
	for _, pindex := range pindexes {
		if pindex.IndexType == "bleve" &&
			pindex.IndexName == indexName &&
			(indexUUID == "" || pindex.IndexUUID == indexUUID) {
			bindex, ok := pindex.Impl.(bleve.Index)
			if ok && bindex != nil {
				alias.Add(bindex)
			}
		}
	}

	return alias, nil
}

// ---------------------------------------------------

type ListIndexHandler struct {
	mgr *Manager
}

func NewListIndexHandler(mgr *Manager) *ListIndexHandler {
	return &ListIndexHandler{mgr: mgr}
}

func (h *ListIndexHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexDefs, _, err := h.mgr.GetIndexDefs(false)
	if err != nil {
		showError(w, req, "could not retrieve index defs", 500)
		return
	}

	rv := struct {
		Status    string     `json:"status"`
		IndexDefs *IndexDefs `json:"indexDefs"`
	}{
		Status:    "ok",
		IndexDefs: indexDefs,
	}
	mustEncode(w, rv)
}

// ---------------------------------------------------

type GetIndexHandler struct {
	mgr *Manager
}

func NewGetIndexHandler(mgr *Manager) *GetIndexHandler {
	return &GetIndexHandler{mgr: mgr}
}

func (h *GetIndexHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexName := indexNameLookup(req)
	if indexName == "" {
		showError(w, req, "index name is required", 400)
		return
	}

	_, indexDefsByName, err := h.mgr.GetIndexDefs(false)
	if err != nil {
		showError(w, req, "could not retrieve index defs", 500)
		return
	}

	indexDef, exists := indexDefsByName[indexName]
	if !exists || indexDef == nil {
		showError(w, req, "not an index", 400)
		return
	}

	indexUUID := req.FormValue("indexUUID")
	if indexUUID != "" && indexUUID != indexDef.UUID {
		showError(w, req, "wrong index UUID", 400)
		return
	}

	m := map[string]interface{}{}
	if indexDef.Schema != "" {
		if err := json.Unmarshal([]byte(indexDef.Schema), &m); err != nil {
			showError(w, req, "could not unmarshal mapping", 500)
			return
		}
	}

	rv := struct {
		Status       string                 `json:"status"`
		IndexDef     *IndexDef              `json:"indexDef"`
		IndexMapping map[string]interface{} `json:"indexMapping"`
	}{
		Status:       "ok",
		IndexDef:     indexDef,
		IndexMapping: m,
	}
	mustEncode(w, rv)
}

// ---------------------------------------------------

type CountHandler struct {
	mgr *Manager
}

func NewCountHandler(mgr *Manager) *CountHandler {
	return &CountHandler{mgr: mgr}
}

func (h *CountHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexName := indexNameLookup(req)
	if indexName == "" {
		showError(w, req, "index name is required", 400)
		return
	}

	indexUUID := req.FormValue("indexUUID")

	alias, err := indexAlias(h.mgr, indexName, indexUUID)
	if err != nil {
		showError(w, req, fmt.Sprintf("index alias: %v", err), 500)
		return
	}

	docCount, err := alias.DocCount()
	if err != nil {
		showError(w, req, fmt.Sprintf("error counting docs: %v", err), 500)
		return
	}

	rv := struct {
		Status string `json:"status"`
		Count  uint64 `json:"count"`
	}{
		Status: "ok",
		Count:  docCount,
	}
	mustEncode(w, rv)
}

// ---------------------------------------------------

type SearchHandler struct {
	mgr *Manager
}

func NewSearchHandler(mgr *Manager) *SearchHandler {
	return &SearchHandler{mgr: mgr}
}

func (h *SearchHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexName := indexNameLookup(req)
	if indexName == "" {
		showError(w, req, "index name is required", 400)
		return
	}

	indexUUID := req.FormValue("indexUUID")

	log.Printf("search request: %s", indexName)

	alias, err := indexAlias(h.mgr, indexName, indexUUID)
	if err != nil {
		showError(w, req, fmt.Sprintf("index alias: %v", err), 500)
		return
	}

	// read the request body
	requestBody, err := ioutil.ReadAll(req.Body)
	if err != nil {
		showError(w, req, fmt.Sprintf("error reading request body: %v", err), 400)
		return
	}

	log.Printf("search request body: %s", requestBody)

	// parse the request
	var searchRequest bleve.SearchRequest
	err = json.Unmarshal(requestBody, &searchRequest)
	if err != nil {
		showError(w, req, fmt.Sprintf("error parsing query: %v", err), 400)
		return
	}

	log.Printf("search parsed request %#v", searchRequest)

	// varlidate the query
	err = searchRequest.Query.Validate()
	if err != nil {
		showError(w, req, fmt.Sprintf("error validating query: %v", err), 400)
		return
	}

	// execute the query
	searchResponse, err := alias.Search(&searchRequest)
	if err != nil {
		showError(w, req, fmt.Sprintf("error executing query: %v", err), 500)
		return
	}

	// encode the response
	mustEncode(w, searchResponse)
}

// ---------------------------------------------------

type FeedStatsHandler struct {
	mgr *Manager
}

func NewFeedStatsHandler(mgr *Manager) *FeedStatsHandler {
	return &FeedStatsHandler{mgr: mgr}
}

func (h *FeedStatsHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	feeds, _ := h.mgr.CurrentMaps()
	w.Write([]byte("[\n"))
	first := true
	feedNames := make([]string, 0, len(feeds))
	for feedName := range feeds {
		feedNames = append(feedNames, feedName)
	}
	sort.Strings(feedNames)
	for _, feedName := range feedNames {
		if !first {
			w.Write([]byte(",\n"))
		}
		first = false
		w.Write([]byte(fmt.Sprintf("  {\"feedName\":\"%s\",\"stats\":", feedName)))
		feeds[feedName].Stats(w)
		w.Write([]byte("}\n"))
	}
	w.Write([]byte("]\n"))
}
<|endoftext|>