<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package backups_test

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"

	gc "launchpad.net/gocheck"

	"github.com/juju/juju/state/backups"
	"github.com/juju/juju/testing"
)

var _ = gc.Suite(&legacySuite{})

type legacySuite struct {
	testing.BaseSuite
	cwd       string
	testFiles []string
}

func (s *legacySuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	s.cwd = c.MkDir()
}

func (s *legacySuite) createTestFiles(c *gc.C) {
	tarDirE := path.Join(s.cwd, "TarDirectoryEmpty")
	err := os.Mkdir(tarDirE, os.FileMode(0755))
	c.Check(err, gc.IsNil)

	tarDirP := path.Join(s.cwd, "TarDirectoryPopulated")
	err = os.Mkdir(tarDirP, os.FileMode(0755))
	c.Check(err, gc.IsNil)

	tarSubFile1 := path.Join(tarDirP, "TarSubFile1")
	tarSubFile1Handle, err := os.Create(tarSubFile1)
	c.Check(err, gc.IsNil)
	tarSubFile1Handle.WriteString("TarSubFile1")
	tarSubFile1Handle.Close()

	tarSubDir := path.Join(tarDirP, "TarDirectoryPopulatedSubDirectory")
	err = os.Mkdir(tarSubDir, os.FileMode(0755))
	c.Check(err, gc.IsNil)

	tarFile1 := path.Join(s.cwd, "TarFile1")
	tarFile1Handle, err := os.Create(tarFile1)
	c.Check(err, gc.IsNil)
	tarFile1Handle.WriteString("TarFile1")
	tarFile1Handle.Close()

	tarFile2 := path.Join(s.cwd, "TarFile2")
	tarFile2Handle, err := os.Create(tarFile2)
	c.Check(err, gc.IsNil)
	tarFile2Handle.WriteString("TarFile2")
	tarFile2Handle.Close()
	s.testFiles = []string{tarDirE, tarDirP, tarFile1, tarFile2}

}

type expectedTarContents struct {
	Name string
	Body string
}

var testExpectedTarContents = []expectedTarContents{
	{"TarDirectoryEmpty", ""},
	{"TarDirectoryPopulated", ""},
	{"TarDirectoryPopulated/TarSubFile1", "TarSubFile1"},
	{"TarDirectoryPopulated/TarDirectoryPopulatedSubDirectory", ""},
	{"TarFile1", "TarFile1"},
	{"TarFile2", "TarFile2"},
}

// Assert thar contents checks that the tar[.gz] file provided contains the
// Expected files
// expectedContents: is a slice of the filenames with relative paths that are
// expected to be on the tar file
// tarFile: is the path of the file to be checked
func (s *legacySuite) assertTarContents(c *gc.C, expectedContents []expectedTarContents,
	tarFile string,
	compressed bool) {
	f, err := os.Open(tarFile)
	c.Assert(err, gc.IsNil)
	defer f.Close()
	var r io.Reader = f
	if compressed {
		r, err = gzip.NewReader(r)
		c.Assert(err, gc.IsNil)
	}

	tr := tar.NewReader(r)

	tarContents := make(map[string]string)
	// Iterate through the files in the archive.
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		c.Assert(err, gc.IsNil)
		buf, err := ioutil.ReadAll(tr)
		c.Assert(err, gc.IsNil)
		tarContents[hdr.Name] = string(buf)
	}
	for _, expectedContent := range expectedContents {
		fullExpectedContent := strings.TrimPrefix(expectedContent.Name, string(os.PathSeparator))
		body, ok := tarContents[fullExpectedContent]
		c.Log(tarContents)
		c.Log(expectedContents)
		c.Log(fmt.Sprintf("checking for presence of %q on tar file", fullExpectedContent))
		c.Assert(ok, gc.Equals, true)
		if expectedContent.Body != "" {
			c.Log("Also checking the file contents")
			c.Assert(body, gc.Equals, expectedContent.Body)
		}
	}

}

func (s *legacySuite) TestBackup(c *gc.C) {
	s.createTestFiles(c)

	s.PatchValue(backups.GetMongodumpPath, func() (string, error) {
		return "bogusmongodump", nil
	})
	s.PatchValue(backups.GetFilesToBackup, func() ([]string, error) {
		return s.testFiles, nil
	})
	ranCommand := false
	s.PatchValue(backups.RunCommand, func(command string, args ...string) error {
		ranCommand = true
		return nil
	})

	bkpFile, shaSum, err := backups.Backup("boguspassword", "bogus-user", s.cwd, "localhost:8080")
	c.Check(err, gc.IsNil)
	c.Assert(ranCommand, gc.Equals, true)

	// It is important that the filename uses non-special characters
	// only because it is returned in a header (unencoded) by the
	// backup API call. This also avoids compatibility problems with
	// client side filename conventions.
	c.Check(bkpFile, gc.Matches, `^[a-z0-9_.-]+$`)

	fileShaSum := shaSumFile(c, path.Join(s.cwd, bkpFile))
	c.Assert(shaSum, gc.Equals, fileShaSum)

	bkpExpectedContents := []expectedTarContents{
		{"juju-backup", ""},
		{"juju-backup/dump", ""},
		{"juju-backup/root.tar", ""},
	}
	s.assertTarContents(c, bkpExpectedContents, path.Join(s.cwd, bkpFile), true)
}

func (s *legacySuite) TestStorageName(c *gc.C) {
	c.Check(backups.StorageName("foo"), gc.Equals, "/backups/foo")
	c.Check(backups.StorageName("/foo/bar"), gc.Equals, "/backups/bar")
	c.Check(backups.StorageName("foo/bar"), gc.Equals, "/backups/bar")
}
<commit_msg>Fix formatting in backups tests.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package backups_test

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"

	gc "launchpad.net/gocheck"

	"github.com/juju/juju/state/backups"
	"github.com/juju/juju/testing"
)

var _ = gc.Suite(&legacySuite{})

type legacySuite struct {
	testing.BaseSuite
	cwd       string
	testFiles []string
}

func (s *legacySuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	s.cwd = c.MkDir()
}

func (s *legacySuite) createTestFiles(c *gc.C) {
	tarDirE := path.Join(s.cwd, "TarDirectoryEmpty")
	err := os.Mkdir(tarDirE, os.FileMode(0755))
	c.Check(err, gc.IsNil)

	tarDirP := path.Join(s.cwd, "TarDirectoryPopulated")
	err = os.Mkdir(tarDirP, os.FileMode(0755))
	c.Check(err, gc.IsNil)

	tarSubFile1 := path.Join(tarDirP, "TarSubFile1")
	tarSubFile1Handle, err := os.Create(tarSubFile1)
	c.Check(err, gc.IsNil)
	tarSubFile1Handle.WriteString("TarSubFile1")
	tarSubFile1Handle.Close()

	tarSubDir := path.Join(tarDirP, "TarDirectoryPopulatedSubDirectory")
	err = os.Mkdir(tarSubDir, os.FileMode(0755))
	c.Check(err, gc.IsNil)

	tarFile1 := path.Join(s.cwd, "TarFile1")
	tarFile1Handle, err := os.Create(tarFile1)
	c.Check(err, gc.IsNil)
	tarFile1Handle.WriteString("TarFile1")
	tarFile1Handle.Close()

	tarFile2 := path.Join(s.cwd, "TarFile2")
	tarFile2Handle, err := os.Create(tarFile2)
	c.Check(err, gc.IsNil)
	tarFile2Handle.WriteString("TarFile2")
	tarFile2Handle.Close()
	s.testFiles = []string{tarDirE, tarDirP, tarFile1, tarFile2}

}

type expectedTarContents struct {
	Name string
	Body string
}

var testExpectedTarContents = []expectedTarContents{
	{"TarDirectoryEmpty", ""},
	{"TarDirectoryPopulated", ""},
	{"TarDirectoryPopulated/TarSubFile1", "TarSubFile1"},
	{"TarDirectoryPopulated/TarDirectoryPopulatedSubDirectory", ""},
	{"TarFile1", "TarFile1"},
	{"TarFile2", "TarFile2"},
}

// Assert thar contents checks that the tar[.gz] file provided contains the
// Expected files
// expectedContents: is a slice of the filenames with relative paths that are
// expected to be on the tar file
// tarFile: is the path of the file to be checked
func (s *legacySuite) assertTarContents(
	c *gc.C, expected []expectedTarContents, tarFile string, compressed bool,
) {
	f, err := os.Open(tarFile)
	c.Assert(err, gc.IsNil)
	defer f.Close()
	var r io.Reader = f
	if compressed {
		r, err = gzip.NewReader(r)
		c.Assert(err, gc.IsNil)
	}

	tr := tar.NewReader(r)

	tarContents := make(map[string]string)
	// Iterate through the files in the archive.
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		c.Assert(err, gc.IsNil)
		buf, err := ioutil.ReadAll(tr)
		c.Assert(err, gc.IsNil)
		tarContents[hdr.Name] = string(buf)
	}
	for _, expectedContent := range expected {
		fullExpectedContent := strings.TrimPrefix(expectedContent.Name, string(os.PathSeparator))
		body, ok := tarContents[fullExpectedContent]
		c.Log(tarContents)
		c.Log(expected)
		c.Log(fmt.Sprintf("checking for presence of %q on tar file", fullExpectedContent))
		c.Assert(ok, gc.Equals, true)
		if expectedContent.Body != "" {
			c.Log("Also checking the file contents")
			c.Assert(body, gc.Equals, expectedContent.Body)
		}
	}

}

func (s *legacySuite) TestBackup(c *gc.C) {
	s.createTestFiles(c)

	s.PatchValue(backups.GetMongodumpPath, func() (string, error) {
		return "bogusmongodump", nil
	})
	s.PatchValue(backups.GetFilesToBackup, func() ([]string, error) {
		return s.testFiles, nil
	})
	ranCommand := false
	s.PatchValue(backups.RunCommand, func(command string, args ...string) error {
		ranCommand = true
		return nil
	})

	bkpFile, shaSum, err := backups.Backup("boguspassword", "bogus-user", s.cwd, "localhost:8080")
	c.Check(err, gc.IsNil)
	c.Assert(ranCommand, gc.Equals, true)

	// It is important that the filename uses non-special characters
	// only because it is returned in a header (unencoded) by the
	// backup API call. This also avoids compatibility problems with
	// client side filename conventions.
	c.Check(bkpFile, gc.Matches, `^[a-z0-9_.-]+$`)

	fileShaSum := shaSumFile(c, path.Join(s.cwd, bkpFile))
	c.Assert(shaSum, gc.Equals, fileShaSum)

	bkpExpectedContents := []expectedTarContents{
		{"juju-backup", ""},
		{"juju-backup/dump", ""},
		{"juju-backup/root.tar", ""},
	}
	s.assertTarContents(c, bkpExpectedContents, path.Join(s.cwd, bkpFile), true)
}

func (s *legacySuite) TestStorageName(c *gc.C) {
	c.Check(backups.StorageName("foo"), gc.Equals, "/backups/foo")
	c.Check(backups.StorageName("/foo/bar"), gc.Equals, "/backups/bar")
	c.Check(backups.StorageName("foo/bar"), gc.Equals, "/backups/bar")
}
<|endoftext|>


<commit_before>package giniapi

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"time"
)

// Timing struct
type Timing struct {
	Upload     time.Duration
	Processing time.Duration
}

// Total returns the summarized timings of upload and processing
func (t *Timing) Total() time.Duration {
	return t.Upload + t.Processing
}

// Page describes a documents pages
type Page struct {
	Images     map[string]string `json:"images"`
	PageNumber int               `json:"pageNumber"`
}

// Links contains the links to a documents resources
type Links struct {
	Document    string `json:"document"`
	Extractions string `json:"extractions"`
	Layout      string `json:"layout"`
	Processed   string `json:"processed"`
}

// Document contains all informations about a single document
type Document struct {
	Timing               `json:"-"`
	client               *APIClient
	Owner                string `json:"-"`
	Links                Links  `json:"_links"`
	CreationDate         int    `json:"creationDate"`
	ID                   string `json:"id"`
	Name                 string `json:"name"`
	Origin               string `json:"origin"`
	PageCount            int    `json:"pageCount"`
	Pages                []Page `json:"pages"`
	Progress             string `json:"progress"`
	SourceClassification string `json:"sourceClassification"`
}

// DocumentSet is a list of documents with the total count
type DocumentSet struct {
	TotalCount int         `json:"totalCount"`
	Documents  []*Document `json:"documents"`
}

// String representaion of a document
func (d *Document) String() string {
	return fmt.Sprintf(d.ID)
}

// Poll the progress state of a document and return nil when the processing
// has completed (successful or failed). On timeout return error
func (d *Document) Poll(timeout time.Duration) error {
	start := time.Now()
	defer func() { d.Timing.Processing = time.Since(start) }()

	docProgress := make(chan *Document, 1)
	quit := make(chan bool, 1)

	go func() {
		for {
			select {
			case <-quit:
				return
			default:
				doc, err := d.client.Get(d.Links.Document, d.Owner)
				if err != nil {
					return
				}
				if doc.Progress == "COMPLETED" || doc.Progress == "ERROR" {
					docProgress <- doc
					return
				}
			}
		}
	}()

	select {
	case doc := <-docProgress:
		if doc == nil {
			return newHTTPError(ErrDocumentProcessing, "", nil, nil)
		}
		*d = *doc
		return nil
	case <-time.After(timeout):
		quit <- true
		return newHTTPError(fmt.Sprintf("%s after %f", ErrDocumentTimeout, timeout.Seconds()), "", nil, nil)
	}
}

// Update document struct from self-contained document link
func (d *Document) Update() error {
	newDoc, err := d.client.Get(d.Links.Document, d.Owner)
	if err != nil {
		return err
	}
	*d = *newDoc
	return nil
}

// Delete a document
func (d *Document) Delete() error {
	resp, err := d.client.makeAPIRequest("DELETE", d.Links.Document, nil, nil, d.Owner)

	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusNoContent {
		return newHTTPError(ErrDocumentDelete, d.ID, err, resp)
	}

	return nil
}

// ErrorReport creates a bug report in Gini's bugtracking system. It's a convinience way
// to help Gini learn from difficult documents
func (d *Document) ErrorReport(summary string, description string) error {
	params := map[string]interface{}{
		"summary":     summary,
		"description": description,
	}

	u := encodeURLParams(fmt.Sprintf("%s/errorreport", d.Links.Document), params)

	resp, err := d.client.makeAPIRequest("POST", u, nil, nil, d.Owner)

	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return newHTTPError(ErrDocumentReport, d.ID, err, resp)
	}

	return nil
}

// GetLayout returns the JSON representation of a documents layout parsed as
// Layout struct
func (d *Document) GetLayout() (*Layout, error) {
	var layout Layout

	resp, err := d.client.makeAPIRequest("GET", d.Links.Layout, nil, nil, "")

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, newHTTPError(ErrDocumentLayout, d.ID, err, resp)
	}

	if err := json.NewDecoder(resp.Body).Decode(&layout); err != nil {
		return nil, err
	}

	return &layout, nil
}

// GetExtractions returns a documents extractions in a Extractions struct
func (d *Document) GetExtractions(incubator bool) (*Extractions, error) {
	var extractions Extractions
	var headers map[string]string

	if incubator {
		headers = map[string]string{
			"Accept": "application/vnd.gini.incubator+json",
		}
	}

	resp, err := d.client.makeAPIRequest("GET", d.Links.Extractions, nil, headers, d.Owner)

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, newHTTPError(ErrDocumentExtractions, d.ID, err, resp)
	}

	if err := json.NewDecoder(resp.Body).Decode(&extractions); err != nil {
		return nil, err
	}

	return &extractions, nil
}

// GetProcessed returns a byte array of the processed (rectified, optimized) document
func (d *Document) GetProcessed() ([]byte, error) {
	headers := map[string]string{
		"Accept": "application/octet-stream",
	}

	resp, err := d.client.makeAPIRequest("GET", d.Links.Processed, nil, headers, d.Owner)
	defer resp.Body.Close()

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, newHTTPError(ErrDocumentProcessed, d.ID, err, resp)
	}

	buf := new(bytes.Buffer)
	_, err = buf.ReadFrom(resp.Body)

	if err != nil {
		return nil, newHTTPError(ErrDocumentProcessed, d.ID, err, resp)
	}

	return buf.Bytes(), nil
}

// SubmitFeedback submits feedback from map
func (d *Document) SubmitFeedback(feedback map[string]map[string]interface{}) error {
	feedbackMap := map[string]map[string]map[string]interface{}{
		"feedback": feedback,
	}

	feedbackBody, err := json.Marshal(feedbackMap)
	if err != nil {
		return err
	}

	resp, err := d.client.makeAPIRequest("PUT", d.Links.Extractions, bytes.NewReader(feedbackBody), nil, d.Owner)
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusNoContent {
		return newHTTPError(ErrDocumentFeedback, d.ID, err, resp)
	}

	return nil
}
<commit_msg>Fix missing upload duration. Was overwritten by d.Poll()<commit_after>package giniapi

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"time"
)

// Timing struct
type Timing struct {
	Upload     time.Duration
	Processing time.Duration
}

// Total returns the summarized timings of upload and processing
func (t *Timing) Total() time.Duration {
	return t.Upload + t.Processing
}

// Page describes a documents pages
type Page struct {
	Images     map[string]string `json:"images"`
	PageNumber int               `json:"pageNumber"`
}

// Links contains the links to a documents resources
type Links struct {
	Document    string `json:"document"`
	Extractions string `json:"extractions"`
	Layout      string `json:"layout"`
	Processed   string `json:"processed"`
}

// Document contains all informations about a single document
type Document struct {
	Timing               `json:"-"`
	client               *APIClient
	Owner                string `json:"-"`
	Links                Links  `json:"_links"`
	CreationDate         int    `json:"creationDate"`
	ID                   string `json:"id"`
	Name                 string `json:"name"`
	Origin               string `json:"origin"`
	PageCount            int    `json:"pageCount"`
	Pages                []Page `json:"pages"`
	Progress             string `json:"progress"`
	SourceClassification string `json:"sourceClassification"`
}

// DocumentSet is a list of documents with the total count
type DocumentSet struct {
	TotalCount int         `json:"totalCount"`
	Documents  []*Document `json:"documents"`
}

// String representaion of a document
func (d *Document) String() string {
	return fmt.Sprintf(d.ID)
}

// Poll the progress state of a document and return nil when the processing
// has completed (successful or failed). On timeout return error
func (d *Document) Poll(timeout time.Duration) error {
	// store upload duration. Will be overwritten otherwise
	uploadDuration := d.Timing.Upload

	start := time.Now()
	defer func() { d.Timing.Processing = time.Since(start) }()

	docProgress := make(chan *Document, 1)
	quit := make(chan bool, 1)

	go func() {
		for {
			select {
			case <-quit:
				return
			default:
				doc, err := d.client.Get(d.Links.Document, d.Owner)
				if err != nil {
					return
				}
				if doc.Progress == "COMPLETED" || doc.Progress == "ERROR" {
					docProgress <- doc
					return
				}
			}
		}
	}()

	select {
	case doc := <-docProgress:
		if doc == nil {
			return newHTTPError(ErrDocumentProcessing, "", nil, nil)
		}
		*d = *doc

		// restore upload duration
		d.Timing.Upload = uploadDuration

		return nil
	case <-time.After(timeout):
		quit <- true
		return newHTTPError(fmt.Sprintf("%s after %f", ErrDocumentTimeout, timeout.Seconds()), "", nil, nil)
	}
}

// Update document struct from self-contained document link
func (d *Document) Update() error {
	newDoc, err := d.client.Get(d.Links.Document, d.Owner)
	if err != nil {
		return err
	}
	*d = *newDoc
	return nil
}

// Delete a document
func (d *Document) Delete() error {
	resp, err := d.client.makeAPIRequest("DELETE", d.Links.Document, nil, nil, d.Owner)

	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusNoContent {
		return newHTTPError(ErrDocumentDelete, d.ID, err, resp)
	}

	return nil
}

// ErrorReport creates a bug report in Gini's bugtracking system. It's a convinience way
// to help Gini learn from difficult documents
func (d *Document) ErrorReport(summary string, description string) error {
	params := map[string]interface{}{
		"summary":     summary,
		"description": description,
	}

	u := encodeURLParams(fmt.Sprintf("%s/errorreport", d.Links.Document), params)

	resp, err := d.client.makeAPIRequest("POST", u, nil, nil, d.Owner)

	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return newHTTPError(ErrDocumentReport, d.ID, err, resp)
	}

	return nil
}

// GetLayout returns the JSON representation of a documents layout parsed as
// Layout struct
func (d *Document) GetLayout() (*Layout, error) {
	var layout Layout

	resp, err := d.client.makeAPIRequest("GET", d.Links.Layout, nil, nil, "")

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, newHTTPError(ErrDocumentLayout, d.ID, err, resp)
	}

	if err := json.NewDecoder(resp.Body).Decode(&layout); err != nil {
		return nil, err
	}

	return &layout, nil
}

// GetExtractions returns a documents extractions in a Extractions struct
func (d *Document) GetExtractions(incubator bool) (*Extractions, error) {
	var extractions Extractions
	var headers map[string]string

	if incubator {
		headers = map[string]string{
			"Accept": "application/vnd.gini.incubator+json",
		}
	}

	resp, err := d.client.makeAPIRequest("GET", d.Links.Extractions, nil, headers, d.Owner)

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, newHTTPError(ErrDocumentExtractions, d.ID, err, resp)
	}

	if err := json.NewDecoder(resp.Body).Decode(&extractions); err != nil {
		return nil, err
	}

	return &extractions, nil
}

// GetProcessed returns a byte array of the processed (rectified, optimized) document
func (d *Document) GetProcessed() ([]byte, error) {
	headers := map[string]string{
		"Accept": "application/octet-stream",
	}

	resp, err := d.client.makeAPIRequest("GET", d.Links.Processed, nil, headers, d.Owner)
	defer resp.Body.Close()

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, newHTTPError(ErrDocumentProcessed, d.ID, err, resp)
	}

	buf := new(bytes.Buffer)
	_, err = buf.ReadFrom(resp.Body)

	if err != nil {
		return nil, newHTTPError(ErrDocumentProcessed, d.ID, err, resp)
	}

	return buf.Bytes(), nil
}

// SubmitFeedback submits feedback from map
func (d *Document) SubmitFeedback(feedback map[string]map[string]interface{}) error {
	feedbackMap := map[string]map[string]map[string]interface{}{
		"feedback": feedback,
	}

	feedbackBody, err := json.Marshal(feedbackMap)
	if err != nil {
		return err
	}

	resp, err := d.client.makeAPIRequest("PUT", d.Links.Extractions, bytes.NewReader(feedbackBody), nil, d.Owner)
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusNoContent {
		return newHTTPError(ErrDocumentFeedback, d.ID, err, resp)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
)

type linkUrls struct {
	Self string `json:"self"`
	Git  string `json:"git"`
	HTML string `json:"html"`
}

type dirContent struct {
	Type        string   `json:"type"`
	Size        uint     `json:"size"`
	Name        string   `json:"name"`
	Path        string   `json:"path"`
	SHA         string   `json:"sha"`
	URL         string   `json:"url"`
	GitURL      string   `json:"git_url"`
	HTMLURL     string   `json:"html_url"`
	DownloadURL string   `json:"download_url"`
	Links       linkUrls `json:"_links"`
}

type blobContent struct {
	Content  string `json:"content"`
	Encoding string `json:"encoding"`
	URL      string `json:"url"`
	SHA      string `json:"sha"`
	Size     uint   `json:"size"`
}

const maxRequests = 10
const repoOwnerName = "KhronosGroup"

var specRepoName = "OpenGL-Registry"
var specRepoFolder = "xml"
var specRegexp = regexp.MustCompile(`^(gl|glx|wgl)\.xml$`)
var eglRepoName = "EGL-Registry"
var eglRepoFolder = "api"
var eglRegexp = regexp.MustCompile(`^(egl)\.xml$`)
var docRepoName = "OpenGL-Refpages"
var docRepoFolders = []string{
	"es1.1",
	"es2.0",
	"es3.0",
	"es3.1",
	"es3",
	"gl2.1",
	"gl4",
}
var docRegexp = regexp.MustCompile(`^[ew]?gl[^u_].*\.xml$`)

func validatedAuthHeader(username string, password string) (string, error) {
	client := &http.Client{}
	req, err := http.NewRequest("GET", "https://api.github.com/user", nil)

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

	autStr := fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%s:%s", username, password))))
	req.Header.Add("Authorization", autStr)
	resp, err := client.Do(req)

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

	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", errors.New("GitHub authorization failed")
	}

	return autStr, nil
}

func download(name string, args []string) {
	flags := flag.NewFlagSet(name, flag.ExitOnError)
	xmlDir := flags.String("d", "xml", "XML directory")
	flags.Parse(args)

	specDir := filepath.Join(*xmlDir, "spec")
	if err := os.MkdirAll(specDir, 0755); err != nil {
		log.Fatalln("error creating specification output directory:", err)
	}

	docDir := filepath.Join(*xmlDir, "doc")
	if err := os.MkdirAll(docDir, 0755); err != nil {
		log.Fatalln("error creating documentation output directory:", err)
	}

	reader := bufio.NewReader(os.Stdin)
	fmt.Print("Enter GitHub username: ")
	input, _ := reader.ReadString('\n')
	username := strings.Trim(input, "\n")
	fmt.Print("Enter GitHub password: ")
	input, _ = reader.ReadString('\n')
	password := strings.Trim(input, "\n")

	authHeader, err := validatedAuthHeader(username, password)

	if err != nil {
		log.Fatalln("error with user authorization:", err)
	}

	err = DownloadGitDir(authHeader, specRepoName, specRepoFolder, specRegexp, specDir)
	if err != nil {
		log.Fatalln("error downloading specification files:", err)
	}

	err = DownloadGitDir(authHeader, eglRepoName, eglRepoFolder, eglRegexp, specDir)
	if err != nil {
		log.Fatalln("error downloading egl file:", err)
	}

	for _, folder := range docRepoFolders {
		if err := DownloadGitDir(authHeader, docRepoName, folder, docRegexp, docDir); err != nil {
			log.Fatalln("error downloading documentation files:", err)
		}
	}
}

// DownloadGitDir reads an Git repo and downloads all the listed (filtered) files.
func DownloadGitDir(authStr string, repoName string, repoFolder string, filter *regexp.Regexp, outDir string) error {
	client := &http.Client{}
	rootDirURL := "https://api.github.com/repos/" + repoOwnerName + "/" + repoName + "/contents/" + repoFolder
	rootBlobURL := "https://api.github.com/repos/" + repoOwnerName + "/" + repoName + "/git/blobs/"
	req, err := http.NewRequest("GET", rootDirURL, nil)
	req.Header.Add("Authorization", authStr)
	req.Header.Add("User-Agent", "go-gl/glow")
	resp, err := client.Do(req)

	if err != nil {
		return err
	}
	defer resp.Body.Close()

	var repoContent []dirContent
	if err := json.NewDecoder(resp.Body).Decode(&repoContent); err != nil {
		return err
	}

	var downloadErr error

	wg := new(sync.WaitGroup)
	c := make(chan int, maxRequests)
	for _, e := range repoContent {
		if filter.MatchString(e.Name) {
			c <- 1
			wg.Add(1)
			file := filepath.Join(outDir, e.Name)
			url := rootBlobURL + e.SHA
			go func(url, file string) {
				defer wg.Done()
				if err := downloadBlob(authStr, url, file); err != nil && downloadErr == nil {
					downloadErr = err
				}
				<-c
			}(url, file)
		}
	}
	wg.Wait()

	return downloadErr
}

func downloadBlob(authStr, url, filePath string) error {
	log.Println("Downloading", filePath)

	client := &http.Client{}
	req, err := http.NewRequest("GET", url, nil)

	if err != nil {
		return err
	}

	req.Header.Add("Authorization", authStr)
	req.Header.Add("User-Agent", "go-gl/glow")
	resp, err := client.Do(req)

	if err != nil {
		return err
	}
	defer resp.Body.Close()

	var blob blobContent
	if err := json.NewDecoder(resp.Body).Decode(&blob); err != nil {
		return err
	}

	data, err := base64.StdEncoding.DecodeString(blob.Content)

	err = ioutil.WriteFile(filePath, data, 0644)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Remove blank lines before error handling.<commit_after>package main

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
)

type linkUrls struct {
	Self string `json:"self"`
	Git  string `json:"git"`
	HTML string `json:"html"`
}

type dirContent struct {
	Type        string   `json:"type"`
	Size        uint     `json:"size"`
	Name        string   `json:"name"`
	Path        string   `json:"path"`
	SHA         string   `json:"sha"`
	URL         string   `json:"url"`
	GitURL      string   `json:"git_url"`
	HTMLURL     string   `json:"html_url"`
	DownloadURL string   `json:"download_url"`
	Links       linkUrls `json:"_links"`
}

type blobContent struct {
	Content  string `json:"content"`
	Encoding string `json:"encoding"`
	URL      string `json:"url"`
	SHA      string `json:"sha"`
	Size     uint   `json:"size"`
}

const maxRequests = 10
const repoOwnerName = "KhronosGroup"

var specRepoName = "OpenGL-Registry"
var specRepoFolder = "xml"
var specRegexp = regexp.MustCompile(`^(gl|glx|wgl)\.xml$`)
var eglRepoName = "EGL-Registry"
var eglRepoFolder = "api"
var eglRegexp = regexp.MustCompile(`^(egl)\.xml$`)
var docRepoName = "OpenGL-Refpages"
var docRepoFolders = []string{
	"es1.1",
	"es2.0",
	"es3.0",
	"es3.1",
	"es3",
	"gl2.1",
	"gl4",
}
var docRegexp = regexp.MustCompile(`^[ew]?gl[^u_].*\.xml$`)

func validatedAuthHeader(username string, password string) (string, error) {
	client := &http.Client{}
	req, err := http.NewRequest("GET", "https://api.github.com/user", nil)
	if err != nil {
		return "", err
	}

	autStr := fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%s:%s", username, password))))
	req.Header.Add("Authorization", autStr)
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}

	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", errors.New("GitHub authorization failed")
	}

	return autStr, nil
}

func download(name string, args []string) {
	flags := flag.NewFlagSet(name, flag.ExitOnError)
	xmlDir := flags.String("d", "xml", "XML directory")
	flags.Parse(args)

	specDir := filepath.Join(*xmlDir, "spec")
	if err := os.MkdirAll(specDir, 0755); err != nil {
		log.Fatalln("error creating specification output directory:", err)
	}

	docDir := filepath.Join(*xmlDir, "doc")
	if err := os.MkdirAll(docDir, 0755); err != nil {
		log.Fatalln("error creating documentation output directory:", err)
	}

	reader := bufio.NewReader(os.Stdin)
	fmt.Print("Enter GitHub username: ")
	input, _ := reader.ReadString('\n')
	username := strings.Trim(input, "\n")
	fmt.Print("Enter GitHub password: ")
	input, _ = reader.ReadString('\n')
	password := strings.Trim(input, "\n")

	authHeader, err := validatedAuthHeader(username, password)
	if err != nil {
		log.Fatalln("error with user authorization:", err)
	}

	err = DownloadGitDir(authHeader, specRepoName, specRepoFolder, specRegexp, specDir)
	if err != nil {
		log.Fatalln("error downloading specification files:", err)
	}

	err = DownloadGitDir(authHeader, eglRepoName, eglRepoFolder, eglRegexp, specDir)
	if err != nil {
		log.Fatalln("error downloading egl file:", err)
	}

	for _, folder := range docRepoFolders {
		if err := DownloadGitDir(authHeader, docRepoName, folder, docRegexp, docDir); err != nil {
			log.Fatalln("error downloading documentation files:", err)
		}
	}
}

// DownloadGitDir reads an Git repo and downloads all the listed (filtered) files.
func DownloadGitDir(authStr string, repoName string, repoFolder string, filter *regexp.Regexp, outDir string) error {
	client := &http.Client{}
	rootDirURL := "https://api.github.com/repos/" + repoOwnerName + "/" + repoName + "/contents/" + repoFolder
	rootBlobURL := "https://api.github.com/repos/" + repoOwnerName + "/" + repoName + "/git/blobs/"
	req, err := http.NewRequest("GET", rootDirURL, nil)
	req.Header.Add("Authorization", authStr)
	req.Header.Add("User-Agent", "go-gl/glow")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	var repoContent []dirContent
	if err := json.NewDecoder(resp.Body).Decode(&repoContent); err != nil {
		return err
	}

	var downloadErr error

	wg := new(sync.WaitGroup)
	c := make(chan int, maxRequests)
	for _, e := range repoContent {
		if filter.MatchString(e.Name) {
			c <- 1
			wg.Add(1)
			file := filepath.Join(outDir, e.Name)
			url := rootBlobURL + e.SHA
			go func(url, file string) {
				defer wg.Done()
				if err := downloadBlob(authStr, url, file); err != nil && downloadErr == nil {
					downloadErr = err
				}
				<-c
			}(url, file)
		}
	}
	wg.Wait()

	return downloadErr
}

func downloadBlob(authStr, url, filePath string) error {
	log.Println("Downloading", filePath)

	client := &http.Client{}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}

	req.Header.Add("Authorization", authStr)
	req.Header.Add("User-Agent", "go-gl/glow")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	var blob blobContent
	if err := json.NewDecoder(resp.Body).Decode(&blob); err != nil {
		return err
	}

	data, err := base64.StdEncoding.DecodeString(blob.Content)

	err = ioutil.WriteFile(filePath, data, 0644)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


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

package mysql

import (
	"context"
	"database/sql"

	"github.com/google/trillian"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/storage"
	"github.com/google/trillian/storage/cache"
	"github.com/google/trillian/types"

	"github.com/golang/glog"
	"github.com/golang/protobuf/proto"
)

const (
	insertMapHeadSQL = `INSERT INTO MapHead(TreeId, MapHeadTimestamp, RootHash, MapRevision, RootSignature, MapperData)
	VALUES(?, ?, ?, ?, ?, ?)`
	selectLatestSignedMapRootSQL = `SELECT MapHeadTimestamp, RootHash, MapRevision, RootSignature, MapperData
		 FROM MapHead WHERE TreeId=?
		 ORDER BY MapHeadTimestamp DESC LIMIT 1`
	selectGetSignedMapRootSQL = `SELECT MapHeadTimestamp, RootHash, MapRevision, RootSignature, MapperData
		 FROM MapHead WHERE TreeId=? AND MapRevision=?`
	insertMapLeafSQL = `INSERT INTO MapLeaf(TreeId, KeyHash, MapRevision, LeafValue) VALUES (?, ?, ?, ?)`
	selectMapLeafSQL = `
 SELECT t1.KeyHash, t1.MapRevision, t1.LeafValue
 FROM MapLeaf t1
 INNER JOIN
 (
	SELECT TreeId, KeyHash, MAX(MapRevision) as maxrev
	FROM MapLeaf t0
	WHERE t0.KeyHash IN (` + placeholderSQL + `) AND
	      t0.TreeId = ? AND t0.MapRevision <= ?
	GROUP BY t0.TreeId, t0.KeyHash
 ) t2
 ON t1.TreeId=t2.TreeId
 AND t1.KeyHash=t2.KeyHash
 AND t1.MapRevision=t2.maxrev`
)

var defaultMapStrata = []int{8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 176}

type mySQLMapStorage struct {
	*mySQLTreeStorage
	admin storage.AdminStorage
}

// NewMapStorage creates a storage.MapStorage instance for the specified MySQL URL.
// It assumes storage.AdminStorage is backed by the same MySQL database as well.
func NewMapStorage(db *sql.DB) storage.MapStorage {
	return &mySQLMapStorage{
		admin:            NewAdminStorage(db),
		mySQLTreeStorage: newTreeStorage(db),
	}
}

func (m *mySQLMapStorage) CheckDatabaseAccessible(ctx context.Context) error {
	return m.db.PingContext(ctx)
}

type readOnlyMapTX struct {
	*sql.Tx
}

func (m *mySQLMapStorage) Snapshot(ctx context.Context) (storage.ReadOnlyMapTX, error) {
	tx, err := m.db.BeginTx(ctx, nil /* opts */)
	if err != nil {
		return nil, err
	}
	return &readOnlyMapTX{tx}, nil
}

func (t *readOnlyMapTX) Close() error {
	if err := t.Rollback(); err != nil && err != sql.ErrTxDone {
		glog.Warningf("Rollback error on Close(): %v", err)
		return err
	}
	return nil
}

func (m *mySQLMapStorage) begin(ctx context.Context, tree *trillian.Tree) (storage.MapTreeTX, error) {
	hasher, err := hashers.NewMapHasher(tree.HashStrategy)
	if err != nil {
		return nil, err
	}

	stCache := cache.NewMapSubtreeCache(defaultMapStrata, tree.TreeId, hasher)
	ttx, err := m.beginTreeTx(ctx, tree, hasher.Size(), stCache)
	if err != nil {
		return nil, err
	}

	mtx := &mapTreeTX{
		treeTX: ttx,
		ms:     m,
	}

	mtx.root, err = mtx.LatestSignedMapRoot(ctx)
	if err != nil && err != storage.ErrTreeNeedsInit {
		return nil, err
	}
	if err == storage.ErrTreeNeedsInit {
		return mtx, err
	}

	if err := mtx.smr.UnmarshalBinary(mtx.root.MapRoot); err != nil {
		return nil, err
	}

	mtx.treeTX.writeRevision = int64(mtx.smr.Revision) + 1
	return mtx, nil
}

func (m *mySQLMapStorage) SnapshotForTree(ctx context.Context, tree *trillian.Tree) (storage.ReadOnlyMapTreeTX, error) {
	return m.begin(ctx, tree)
}

func (m *mySQLMapStorage) ReadWriteTransaction(ctx context.Context, tree *trillian.Tree, f storage.MapTXFunc) error {
	tx, err := m.begin(ctx, tree)
	if tx != nil {
		defer tx.Close()
	}
	if err != nil && err != storage.ErrTreeNeedsInit {
		return err
	}
	if err := f(ctx, tx); err != nil {
		return err
	}
	return tx.Commit()
}

type mapTreeTX struct {
	treeTX
	ms   *mySQLMapStorage
	root trillian.SignedMapRoot
	smr  types.MapRootV1
}

func (m *mapTreeTX) ReadRevision() int64 {
	return int64(m.smr.Revision)
}

func (m *mapTreeTX) WriteRevision() int64 {
	return m.treeTX.writeRevision
}

func (m *mapTreeTX) Set(ctx context.Context, keyHash []byte, value trillian.MapLeaf) error {
	// TODO(al): consider storing some sort of value which represents the group of keys being set in this Tx.
	//           That way, if this attempt partially fails (i.e. because some subset of the in-the-future Merkle
	//           nodes do get written), we can enforce that future map update attempts are a complete replay of
	//           the failed set.
	flatValue, err := proto.Marshal(&value)
	if err != nil {
		return nil
	}

	stmt, err := m.tx.PrepareContext(ctx, insertMapLeafSQL)
	if err != nil {
		return err
	}
	defer stmt.Close()

	_, err = stmt.ExecContext(ctx, m.treeID, keyHash, m.writeRevision, flatValue)
	return err
}

// Get returns a list of map leaves indicated by indexes.
// If an index is not found, no corresponding entry is returned.
// Each MapLeaf.Index is overwritten with the index the leaf was found at.
func (m *mapTreeTX) Get(ctx context.Context, revision int64, indexes [][]byte) ([]trillian.MapLeaf, error) {
	// If no indexes are requested, return an empty set.
	if len(indexes) == 0 {
		return []trillian.MapLeaf{}, nil
	}
	stmt, err := m.ms.getStmt(ctx, selectMapLeafSQL, len(indexes), "?", "?")
	if err != nil {
		return nil, err
	}
	stx := m.tx.StmtContext(ctx, stmt)
	defer stx.Close()

	args := make([]interface{}, 0, len(indexes)+2)
	for _, index := range indexes {
		args = append(args, index)
	}
	args = append(args, m.treeID)
	args = append(args, revision)

	rows, err := stx.QueryContext(ctx, args...)
	// It's possible there are no values for any of these keys yet
	if err == sql.ErrNoRows {
		return nil, nil
	} else if err != nil {
		return nil, err
	}
	defer rows.Close()

	ret := make([]trillian.MapLeaf, 0, len(indexes))
	nr := 0
	er := 0
	for rows.Next() {
		var mapKeyHash []byte
		var mapRevision int64
		var flatData []byte
		err = rows.Scan(&mapKeyHash, &mapRevision, &flatData)
		if err != nil {
			return nil, err
		}
		if len(flatData) == 0 {
			er++
			continue
		}
		var mapLeaf trillian.MapLeaf
		err = proto.Unmarshal(flatData, &mapLeaf)
		if err != nil {
			return nil, err
		}
		mapLeaf.Index = mapKeyHash
		ret = append(ret, mapLeaf)
		nr++
	}
	return ret, nil
}

func (m *mapTreeTX) GetSignedMapRoot(ctx context.Context, revision int64) (trillian.SignedMapRoot, error) {
	var timestamp, mapRevision int64
	var rootHash, rootSignatureBytes []byte
	var mapperMetaBytes []byte

	stmt, err := m.tx.PrepareContext(ctx, selectGetSignedMapRootSQL)
	if err != nil {
		return trillian.SignedMapRoot{}, err
	}
	defer stmt.Close()

	err = stmt.QueryRowContext(ctx, m.treeID, revision).Scan(
		&timestamp, &rootHash, &mapRevision, &rootSignatureBytes, &mapperMetaBytes)
	if err != nil {
		if revision == 0 {
			return trillian.SignedMapRoot{}, storage.ErrTreeNeedsInit
		}
		return trillian.SignedMapRoot{}, err
	}
	return m.signedMapRoot(timestamp, mapRevision, rootHash, rootSignatureBytes, mapperMetaBytes)
}

func (m *mapTreeTX) LatestSignedMapRoot(ctx context.Context) (trillian.SignedMapRoot, error) {
	var timestamp, mapRevision int64
	var rootHash, rootSignatureBytes []byte
	var mapperMetaBytes []byte

	stmt, err := m.tx.PrepareContext(ctx, selectLatestSignedMapRootSQL)
	if err != nil {
		return trillian.SignedMapRoot{}, err
	}
	defer stmt.Close()

	err = stmt.QueryRowContext(ctx, m.treeID).Scan(
		&timestamp, &rootHash, &mapRevision, &rootSignatureBytes, &mapperMetaBytes)

	// It's possible there are no roots for this tree yet
	if err == sql.ErrNoRows {
		return trillian.SignedMapRoot{}, storage.ErrTreeNeedsInit
	} else if err != nil {
		return trillian.SignedMapRoot{}, err
	}
	return m.signedMapRoot(timestamp, mapRevision, rootHash, rootSignatureBytes, mapperMetaBytes)
}

func (m *mapTreeTX) signedMapRoot(timestamp, mapRevision int64, rootHash, rootSignature, mapperMeta []byte) (trillian.SignedMapRoot, error) {
	mapRoot, err := (&types.MapRootV1{
		RootHash:       rootHash,
		TimestampNanos: uint64(timestamp),
		Revision:       uint64(mapRevision),
		Metadata:       mapperMeta,
	}).MarshalBinary()
	if err != nil {
		return trillian.SignedMapRoot{}, err
	}

	return trillian.SignedMapRoot{
		MapRoot:   mapRoot,
		Signature: rootSignature,
	}, nil
}

func (m *mapTreeTX) StoreSignedMapRoot(ctx context.Context, root trillian.SignedMapRoot) error {
	var r types.MapRootV1
	if err := r.UnmarshalBinary(root.MapRoot); err != nil {
		return err
	}

	stmt, err := m.tx.PrepareContext(ctx, insertMapHeadSQL)
	if err != nil {
		return err
	}
	defer stmt.Close()

	// TODO(al): store transactionLogHead too
	res, err := stmt.ExecContext(ctx, m.treeID, r.TimestampNanos, r.RootHash, r.Revision, root.Signature, r.Metadata)

	if err != nil {
		glog.Warningf("Failed to store signed map root: %s", err)
	}

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

package mysql

import (
	"context"
	"database/sql"

	"github.com/google/trillian"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/storage"
	"github.com/google/trillian/storage/cache"
	"github.com/google/trillian/types"

	"github.com/golang/glog"
	"github.com/golang/protobuf/proto"
)

const (
	insertMapHeadSQL = `INSERT INTO MapHead(TreeId, MapHeadTimestamp, RootHash, MapRevision, RootSignature, MapperData)
	VALUES(?, ?, ?, ?, ?, ?)`
	selectLatestSignedMapRootSQL = `SELECT MapHeadTimestamp, RootHash, MapRevision, RootSignature, MapperData
		 FROM MapHead WHERE TreeId=?
		 ORDER BY MapHeadTimestamp DESC LIMIT 1`
	selectGetSignedMapRootSQL = `SELECT MapHeadTimestamp, RootHash, MapRevision, RootSignature, MapperData
		 FROM MapHead WHERE TreeId=? AND MapRevision=?`
	insertMapLeafSQL = `INSERT INTO MapLeaf(TreeId, KeyHash, MapRevision, LeafValue) VALUES (?, ?, ?, ?)`
	selectMapLeafSQL = `
 SELECT t1.KeyHash, t1.MapRevision, t1.LeafValue
 FROM MapLeaf t1
 INNER JOIN
 (
	SELECT TreeId, KeyHash, MAX(MapRevision) as maxrev
	FROM MapLeaf t0
	WHERE t0.KeyHash IN (` + placeholderSQL + `) AND
	      t0.TreeId = ? AND t0.MapRevision <= ?
	GROUP BY t0.TreeId, t0.KeyHash
 ) t2
 ON t1.TreeId=t2.TreeId
 AND t1.KeyHash=t2.KeyHash
 AND t1.MapRevision=t2.maxrev`
)

var defaultMapStrata = []int{8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 176}

type mySQLMapStorage struct {
	*mySQLTreeStorage
	admin storage.AdminStorage
}

// NewMapStorage creates a storage.MapStorage instance for the specified MySQL URL.
// It assumes storage.AdminStorage is backed by the same MySQL database as well.
func NewMapStorage(db *sql.DB) storage.MapStorage {
	return &mySQLMapStorage{
		admin:            NewAdminStorage(db),
		mySQLTreeStorage: newTreeStorage(db),
	}
}

func (m *mySQLMapStorage) CheckDatabaseAccessible(ctx context.Context) error {
	return m.db.PingContext(ctx)
}

type readOnlyMapTX struct {
	*sql.Tx
}

func (m *mySQLMapStorage) Snapshot(ctx context.Context) (storage.ReadOnlyMapTX, error) {
	tx, err := m.db.BeginTx(ctx, nil /* opts */)
	if err != nil {
		return nil, err
	}
	return &readOnlyMapTX{tx}, nil
}

func (t *readOnlyMapTX) Close() error {
	if err := t.Rollback(); err != nil && err != sql.ErrTxDone {
		glog.Warningf("Rollback error on Close(): %v", err)
		return err
	}
	return nil
}

func (m *mySQLMapStorage) begin(ctx context.Context, tree *trillian.Tree) (storage.MapTreeTX, error) {
	hasher, err := hashers.NewMapHasher(tree.HashStrategy)
	if err != nil {
		return nil, err
	}

	stCache := cache.NewMapSubtreeCache(defaultMapStrata, tree.TreeId, hasher)
	ttx, err := m.beginTreeTx(ctx, tree, hasher.Size(), stCache)
	if err != nil {
		return nil, err
	}

	mtx := &mapTreeTX{
		treeTX:       ttx,
		ms:           m,
		readRevision: -1,
	}

	root, err := mtx.LatestSignedMapRoot(ctx)
	if err != nil && err != storage.ErrTreeNeedsInit {
		return nil, err
	}
	if err == storage.ErrTreeNeedsInit {
		return mtx, err
	}

	var mr types.MapRootV1
	if err := mr.UnmarshalBinary(root.MapRoot); err != nil {
		return nil, err
	}

	mtx.readRevision = int64(mr.Revision)
	mtx.treeTX.writeRevision = int64(mr.Revision) + 1
	return mtx, nil
}

func (m *mySQLMapStorage) SnapshotForTree(ctx context.Context, tree *trillian.Tree) (storage.ReadOnlyMapTreeTX, error) {
	return m.begin(ctx, tree)
}

func (m *mySQLMapStorage) ReadWriteTransaction(ctx context.Context, tree *trillian.Tree, f storage.MapTXFunc) error {
	tx, err := m.begin(ctx, tree)
	if tx != nil {
		defer tx.Close()
	}
	if err != nil && err != storage.ErrTreeNeedsInit {
		return err
	}
	if err := f(ctx, tx); err != nil {
		return err
	}
	return tx.Commit()
}

type mapTreeTX struct {
	treeTX
	ms           *mySQLMapStorage
	readRevision int64
}

func (m *mapTreeTX) ReadRevision() int64 {
	return int64(m.readRevision)
}

func (m *mapTreeTX) WriteRevision() int64 {
	return m.treeTX.writeRevision
}

func (m *mapTreeTX) Set(ctx context.Context, keyHash []byte, value trillian.MapLeaf) error {
	// TODO(al): consider storing some sort of value which represents the group of keys being set in this Tx.
	//           That way, if this attempt partially fails (i.e. because some subset of the in-the-future Merkle
	//           nodes do get written), we can enforce that future map update attempts are a complete replay of
	//           the failed set.
	flatValue, err := proto.Marshal(&value)
	if err != nil {
		return nil
	}

	stmt, err := m.tx.PrepareContext(ctx, insertMapLeafSQL)
	if err != nil {
		return err
	}
	defer stmt.Close()

	_, err = stmt.ExecContext(ctx, m.treeID, keyHash, m.writeRevision, flatValue)
	return err
}

// Get returns a list of map leaves indicated by indexes.
// If an index is not found, no corresponding entry is returned.
// Each MapLeaf.Index is overwritten with the index the leaf was found at.
func (m *mapTreeTX) Get(ctx context.Context, revision int64, indexes [][]byte) ([]trillian.MapLeaf, error) {
	// If no indexes are requested, return an empty set.
	if len(indexes) == 0 {
		return []trillian.MapLeaf{}, nil
	}
	stmt, err := m.ms.getStmt(ctx, selectMapLeafSQL, len(indexes), "?", "?")
	if err != nil {
		return nil, err
	}
	stx := m.tx.StmtContext(ctx, stmt)
	defer stx.Close()

	args := make([]interface{}, 0, len(indexes)+2)
	for _, index := range indexes {
		args = append(args, index)
	}
	args = append(args, m.treeID)
	args = append(args, revision)

	rows, err := stx.QueryContext(ctx, args...)
	// It's possible there are no values for any of these keys yet
	if err == sql.ErrNoRows {
		return nil, nil
	} else if err != nil {
		return nil, err
	}
	defer rows.Close()

	ret := make([]trillian.MapLeaf, 0, len(indexes))
	nr := 0
	er := 0
	for rows.Next() {
		var mapKeyHash []byte
		var mapRevision int64
		var flatData []byte
		err = rows.Scan(&mapKeyHash, &mapRevision, &flatData)
		if err != nil {
			return nil, err
		}
		if len(flatData) == 0 {
			er++
			continue
		}
		var mapLeaf trillian.MapLeaf
		err = proto.Unmarshal(flatData, &mapLeaf)
		if err != nil {
			return nil, err
		}
		mapLeaf.Index = mapKeyHash
		ret = append(ret, mapLeaf)
		nr++
	}
	return ret, nil
}

func (m *mapTreeTX) GetSignedMapRoot(ctx context.Context, revision int64) (trillian.SignedMapRoot, error) {
	var timestamp, mapRevision int64
	var rootHash, rootSignatureBytes []byte
	var mapperMetaBytes []byte

	stmt, err := m.tx.PrepareContext(ctx, selectGetSignedMapRootSQL)
	if err != nil {
		return trillian.SignedMapRoot{}, err
	}
	defer stmt.Close()

	err = stmt.QueryRowContext(ctx, m.treeID, revision).Scan(
		&timestamp, &rootHash, &mapRevision, &rootSignatureBytes, &mapperMetaBytes)
	if err != nil {
		if revision == 0 {
			return trillian.SignedMapRoot{}, storage.ErrTreeNeedsInit
		}
		return trillian.SignedMapRoot{}, err
	}
	return m.signedMapRoot(timestamp, mapRevision, rootHash, rootSignatureBytes, mapperMetaBytes)
}

func (m *mapTreeTX) LatestSignedMapRoot(ctx context.Context) (trillian.SignedMapRoot, error) {
	var timestamp, mapRevision int64
	var rootHash, rootSignatureBytes []byte
	var mapperMetaBytes []byte

	stmt, err := m.tx.PrepareContext(ctx, selectLatestSignedMapRootSQL)
	if err != nil {
		return trillian.SignedMapRoot{}, err
	}
	defer stmt.Close()

	err = stmt.QueryRowContext(ctx, m.treeID).Scan(
		&timestamp, &rootHash, &mapRevision, &rootSignatureBytes, &mapperMetaBytes)

	// It's possible there are no roots for this tree yet
	if err == sql.ErrNoRows {
		return trillian.SignedMapRoot{}, storage.ErrTreeNeedsInit
	} else if err != nil {
		return trillian.SignedMapRoot{}, err
	}
	return m.signedMapRoot(timestamp, mapRevision, rootHash, rootSignatureBytes, mapperMetaBytes)
}

func (m *mapTreeTX) signedMapRoot(timestamp, mapRevision int64, rootHash, rootSignature, mapperMeta []byte) (trillian.SignedMapRoot, error) {
	mapRoot, err := (&types.MapRootV1{
		RootHash:       rootHash,
		TimestampNanos: uint64(timestamp),
		Revision:       uint64(mapRevision),
		Metadata:       mapperMeta,
	}).MarshalBinary()
	if err != nil {
		return trillian.SignedMapRoot{}, err
	}

	return trillian.SignedMapRoot{
		MapRoot:   mapRoot,
		Signature: rootSignature,
	}, nil
}

func (m *mapTreeTX) StoreSignedMapRoot(ctx context.Context, root trillian.SignedMapRoot) error {
	var r types.MapRootV1
	if err := r.UnmarshalBinary(root.MapRoot); err != nil {
		return err
	}

	stmt, err := m.tx.PrepareContext(ctx, insertMapHeadSQL)
	if err != nil {
		return err
	}
	defer stmt.Close()

	// TODO(al): store transactionLogHead too
	res, err := stmt.ExecContext(ctx, m.treeID, r.TimestampNanos, r.RootHash, r.Revision, root.Signature, r.Metadata)

	if err != nil {
		glog.Warningf("Failed to store signed map root: %s", err)
	}

	return checkResultOkAndRowCountIs(res, err, 1)
}
<|endoftext|>


<commit_before>package postgres

import (
	"bytes"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"strings"
	"time"

	"honnef.co/go/tracer"

	"github.com/jmoiron/sqlx"
	"github.com/opentracing/opentracing-go"
)

// timeRange represents a PostgreSQL tstzrange. Caveat: it only
// supports inclusive ranges.
type timeRange struct {
	Start time.Time
	End   time.Time
}

func (t *timeRange) Scan(src interface{}) error {
	const layout = "2006-01-02 15:04:05.999999-07"

	b := src.([]byte)
	b = b[2:]
	idx := bytes.IndexByte(b, '"')
	t1, err := time.Parse(layout, string(b[:idx]))
	if err != nil {
		return err
	}

	b = b[idx+1:]
	idx = bytes.IndexByte(b, '"')
	b = b[idx+1:]
	idx = bytes.IndexByte(b, '"')
	t2, err := time.Parse(layout, string(string(b[:idx])))
	if err != nil {
		return err
	}
	t.Start = t1
	t.End = t2
	return nil
}

func (t timeRange) Value() (driver.Value, error) {
	const layout = "2006-01-02 15:04:05.999999-07"
	return []byte(fmt.Sprintf(`["%s","%s"]`, t.Start.Format(layout), t.End.Format(layout))), nil
}

type Storage struct {
	db *sqlx.DB
}

func New(db *sql.DB) *Storage {
	return &Storage{db: sqlx.NewDb(db, "postgres")}
}

func (st *Storage) Store(sp tracer.RawSpan) (err error) {
	const upsertSpan = `
INSERT INTO spans (id, trace_id, time, operation_name)
VALUES ($1, $2, $3, $4)
ON CONFLICT (id) DO
  UPDATE SET
    time = $3,
    operation_name = $4`
	const insertTag = `INSERT INTO tags (span_id, trace_id, key, value) VALUES ($1, $2, $3, $4)`
	const insertLog = `INSERT INTO tags (span_id, trace_id, key, value, time) VALUES ($1, $2, $3, $4, $5)`
	const insertParentRelation = `INSERT INTO relations (span1_id, span2_id, kind) VALUES ($1, $2, 'parent')`
	const insertParentSpan = `INSERT INTO spans (id, trace_id, time, operation_name) VALUES ($1, $2, $3, '') ON CONFLICT (id) DO NOTHING`

	tx, err := st.db.Begin()
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			_ = tx.Rollback()
			return
		}
		err = tx.Commit()
	}()

	_, err = tx.Exec(upsertSpan,
		int64(sp.SpanID), int64(sp.TraceID), timeRange{sp.StartTime, sp.FinishTime}, sp.OperationName)
	if err != nil {
		return err
	}

	if sp.ParentID != 0 {
		_, err = tx.Exec(insertParentSpan,
			int64(sp.ParentID), int64(sp.TraceID), timeRange{time.Time{}, time.Time{}})
		if err != nil {
			return err
		}
		_, err = tx.Exec(insertParentSpan,
			int64(sp.TraceID), int64(sp.TraceID), timeRange{sp.StartTime, sp.FinishTime})
		if err != nil {
			return err
		}

		_, err = tx.Exec(insertParentRelation,
			int64(sp.ParentID), int64(sp.SpanID))
		if err != nil {
			return err
		}
	}

	for k, v := range sp.Tags {
		vs := fmt.Sprintf("%v", v) // XXX
		_, err = tx.Exec(insertTag,
			int64(sp.SpanID), int64(sp.TraceID), k, vs)
		if err != nil {
			return err
		}
	}
	for _, l := range sp.Logs {
		v := fmt.Sprintf("%v", l.Payload) // XXX
		_, err = tx.Exec(insertLog,
			int64(sp.SpanID), int64(sp.TraceID), l.Event, v, l.Timestamp)
		if err != nil {
			return err
		}
	}
	return nil
}

func (st *Storage) TraceWithID(id uint64) (tracer.RawTrace, error) {
	tx, err := st.db.Begin()
	if err != nil {
		return tracer.RawTrace{}, err
	}
	defer tx.Rollback()
	return st.traceWithID(tx, id)
}

func (st *Storage) traceWithID(tx *sql.Tx, id uint64) (tracer.RawTrace, error) {
	const selectTrace = `
SELECT spans.id, spans.trace_id, spans.time, spans.operation_name, tags.key, tags.value, tags.time
FROM spans
  LEFT JOIN tags
    ON spans.id = tags.span_id
WHERE spans.trace_id = $1
ORDER BY
  spans.time ASC,
  spans.id`
	rows, err := tx.Query(selectTrace, int64(id))
	if err != nil {
		return tracer.RawTrace{}, err
	}

	spans, err := scanSpans(rows)
	if err != nil {
		return tracer.RawTrace{}, err
	}
	return tracer.RawTrace{
		TraceID: id,
		Spans:   spans,
	}, nil
}

func scanSpans(rows *sql.Rows) ([]tracer.RawSpan, error) {
	// TODO select parents
	var spans []tracer.RawSpan
	var (
		prevSpanID int64

		spanID        int64
		traceID       int64
		spanTime      timeRange
		operationName string
		tagKey        string
		tagValue      string
		tagTime       *time.Time
	)
	tagTime = new(time.Time)
	var span tracer.RawSpan
	for rows.Next() {
		if err := rows.Scan(&spanID, &traceID, &spanTime, &operationName, &tagKey, &tagValue, &tagTime); err != nil {
			return nil, err
		}
		if spanID != prevSpanID {
			if prevSpanID != 0 {
				spans = append(spans, span)
			}
			prevSpanID = spanID
			span = tracer.RawSpan{
				Tags: map[string]interface{}{},
			}
		}
		span.SpanID = uint64(spanID)
		span.TraceID = uint64(traceID)
		span.StartTime = spanTime.Start
		span.FinishTime = spanTime.End
		span.OperationName = operationName
		if tagKey != "" {
			if tagTime == nil {
				span.Tags[tagKey] = tagValue
			} else {
				span.Logs = append(span.Logs, opentracing.LogData{
					Timestamp: *tagTime,
					Event:     tagKey,
					Payload:   tagValue,
				})
			}
		}
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	if span.SpanID != 0 {
		spans = append(spans, span)
	}
	return spans, nil
}

func (st *Storage) SpanWithID(id uint64) (tracer.RawSpan, error) {
	tx, err := st.db.Begin()
	if err != nil {
		return tracer.RawSpan{}, err
	}
	defer tx.Rollback()
	return st.spanWithID(tx, id)
}

func (st *Storage) spanWithID(tx *sql.Tx, id uint64) (tracer.RawSpan, error) {
	const selectSpan = `
SELECT spans.id, spans.trace_id, spans.time, spans.time, spans.operation_name, tags.key, tags.value, tags.time
FROM spans
  LEFT JOIN tags
    ON spans.id = tags.span_id
WHERE id = $1
LIMIT 1`
	rows, err := tx.Query(selectSpan, int64(id))
	if err != nil {
		return tracer.RawSpan{}, err
	}
	spans, err := scanSpans(rows)
	if err != nil {
		return tracer.RawSpan{}, err
	}
	if len(spans) == 0 {
		return tracer.RawSpan{}, sql.ErrNoRows
	}
	return spans[0], nil
}

func (st *Storage) QueryTraces(q tracer.Query) ([]tracer.RawTrace, error) {
	tx, err := st.db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	var andConds []string
	var andArgs []interface{}
	var orConds []string
	var orArgs []interface{}
	if q.FinishTime.IsZero() {
		q.FinishTime = time.Now()
	}
	for _, tag := range q.AndTags {
		if tag.CheckValue {
			andConds = append(andConds, `(tags.key = ? AND tags.value = ?)`)
			andArgs = append(andArgs, tag.Key, tag.Value)
		} else {
			andConds = append(andConds, `(tags.key = ?)`)
			andArgs = append(andArgs, tag.Key)
		}
	}

	for _, tag := range q.OrTags {
		if tag.CheckValue {
			orConds = append(orConds, `(tags.key = ? AND tags.value = ?)`)
			orArgs = append(orArgs, tag.Key, tag.Value)
		} else {
			orConds = append(orConds, `(tags.key = ?)`)
			orArgs = append(orArgs, tag.Key)
		}
	}

	and := strings.Join(andConds, " AND ")
	or := strings.Join(orConds, " OR ")
	conds := []string{"true"}
	if and != "" {
		conds = append(conds, and)
	}
	if or != "" {
		conds = append(conds, or)
	}

	query := st.db.Rebind(`
SELECT spans.trace_id
FROM spans
WHERE
  EXISTS (
    SELECT 1
    FROM tags
    WHERE
      tags.trace_id = spans.trace_id AND
      ` + strings.Join(conds, " AND ") + `
  ) AND
  ? @> spans.time AND
  (? = '' OR operation_name = ?) AND
  spans.id = spans.trace_id
ORDER BY
  spans.time ASC,
  spans.trace_id
`)
	args := make([]interface{}, 0, len(andArgs)+len(orArgs))
	args = append(args, andArgs...)
	args = append(args, orArgs...)
	args = append(args, timeRange{q.StartTime, q.FinishTime})
	args = append(args, q.OperationName, q.OperationName)

	var ids []int64
	rows, err := st.db.Query(query, args...)
	if err != nil {
		return nil, err
	}
	var id int64
	for rows.Next() {
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	var traces []tracer.RawTrace
	for _, id := range ids {
		trace, err := st.traceWithID(tx, uint64(id))
		if err != nil {
			return nil, err
		}
		traces = append(traces, trace)
	}
	return traces, nil
}
<commit_msg>Ensure that the Queryer and Storer interfaces are implemented<commit_after>package postgres

import (
	"bytes"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"strings"
	"time"

	"honnef.co/go/tracer"

	"github.com/jmoiron/sqlx"
	"github.com/opentracing/opentracing-go"
)

var _ tracer.Queryer = (*Storage)(nil)
var _ tracer.Storer = (*Storage)(nil)

// timeRange represents a PostgreSQL tstzrange. Caveat: it only
// supports inclusive ranges.
type timeRange struct {
	Start time.Time
	End   time.Time
}

func (t *timeRange) Scan(src interface{}) error {
	const layout = "2006-01-02 15:04:05.999999-07"

	b := src.([]byte)
	b = b[2:]
	idx := bytes.IndexByte(b, '"')
	t1, err := time.Parse(layout, string(b[:idx]))
	if err != nil {
		return err
	}

	b = b[idx+1:]
	idx = bytes.IndexByte(b, '"')
	b = b[idx+1:]
	idx = bytes.IndexByte(b, '"')
	t2, err := time.Parse(layout, string(string(b[:idx])))
	if err != nil {
		return err
	}
	t.Start = t1
	t.End = t2
	return nil
}

func (t timeRange) Value() (driver.Value, error) {
	const layout = "2006-01-02 15:04:05.999999-07"
	return []byte(fmt.Sprintf(`["%s","%s"]`, t.Start.Format(layout), t.End.Format(layout))), nil
}

type Storage struct {
	db *sqlx.DB
}

func New(db *sql.DB) *Storage {
	return &Storage{db: sqlx.NewDb(db, "postgres")}
}

func (st *Storage) Store(sp tracer.RawSpan) (err error) {
	const upsertSpan = `
INSERT INTO spans (id, trace_id, time, operation_name)
VALUES ($1, $2, $3, $4)
ON CONFLICT (id) DO
  UPDATE SET
    time = $3,
    operation_name = $4`
	const insertTag = `INSERT INTO tags (span_id, trace_id, key, value) VALUES ($1, $2, $3, $4)`
	const insertLog = `INSERT INTO tags (span_id, trace_id, key, value, time) VALUES ($1, $2, $3, $4, $5)`
	const insertParentRelation = `INSERT INTO relations (span1_id, span2_id, kind) VALUES ($1, $2, 'parent')`
	const insertParentSpan = `INSERT INTO spans (id, trace_id, time, operation_name) VALUES ($1, $2, $3, '') ON CONFLICT (id) DO NOTHING`

	tx, err := st.db.Begin()
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			_ = tx.Rollback()
			return
		}
		err = tx.Commit()
	}()

	_, err = tx.Exec(upsertSpan,
		int64(sp.SpanID), int64(sp.TraceID), timeRange{sp.StartTime, sp.FinishTime}, sp.OperationName)
	if err != nil {
		return err
	}

	if sp.ParentID != 0 {
		_, err = tx.Exec(insertParentSpan,
			int64(sp.ParentID), int64(sp.TraceID), timeRange{time.Time{}, time.Time{}})
		if err != nil {
			return err
		}
		_, err = tx.Exec(insertParentSpan,
			int64(sp.TraceID), int64(sp.TraceID), timeRange{sp.StartTime, sp.FinishTime})
		if err != nil {
			return err
		}

		_, err = tx.Exec(insertParentRelation,
			int64(sp.ParentID), int64(sp.SpanID))
		if err != nil {
			return err
		}
	}

	for k, v := range sp.Tags {
		vs := fmt.Sprintf("%v", v) // XXX
		_, err = tx.Exec(insertTag,
			int64(sp.SpanID), int64(sp.TraceID), k, vs)
		if err != nil {
			return err
		}
	}
	for _, l := range sp.Logs {
		v := fmt.Sprintf("%v", l.Payload) // XXX
		_, err = tx.Exec(insertLog,
			int64(sp.SpanID), int64(sp.TraceID), l.Event, v, l.Timestamp)
		if err != nil {
			return err
		}
	}
	return nil
}

func (st *Storage) TraceWithID(id uint64) (tracer.RawTrace, error) {
	tx, err := st.db.Begin()
	if err != nil {
		return tracer.RawTrace{}, err
	}
	defer tx.Rollback()
	return st.traceWithID(tx, id)
}

func (st *Storage) traceWithID(tx *sql.Tx, id uint64) (tracer.RawTrace, error) {
	const selectTrace = `
SELECT spans.id, spans.trace_id, spans.time, spans.operation_name, tags.key, tags.value, tags.time
FROM spans
  LEFT JOIN tags
    ON spans.id = tags.span_id
WHERE spans.trace_id = $1
ORDER BY
  spans.time ASC,
  spans.id`
	rows, err := tx.Query(selectTrace, int64(id))
	if err != nil {
		return tracer.RawTrace{}, err
	}

	spans, err := scanSpans(rows)
	if err != nil {
		return tracer.RawTrace{}, err
	}
	return tracer.RawTrace{
		TraceID: id,
		Spans:   spans,
	}, nil
}

func scanSpans(rows *sql.Rows) ([]tracer.RawSpan, error) {
	// TODO select parents
	var spans []tracer.RawSpan
	var (
		prevSpanID int64

		spanID        int64
		traceID       int64
		spanTime      timeRange
		operationName string
		tagKey        string
		tagValue      string
		tagTime       *time.Time
	)
	tagTime = new(time.Time)
	var span tracer.RawSpan
	for rows.Next() {
		if err := rows.Scan(&spanID, &traceID, &spanTime, &operationName, &tagKey, &tagValue, &tagTime); err != nil {
			return nil, err
		}
		if spanID != prevSpanID {
			if prevSpanID != 0 {
				spans = append(spans, span)
			}
			prevSpanID = spanID
			span = tracer.RawSpan{
				Tags: map[string]interface{}{},
			}
		}
		span.SpanID = uint64(spanID)
		span.TraceID = uint64(traceID)
		span.StartTime = spanTime.Start
		span.FinishTime = spanTime.End
		span.OperationName = operationName
		if tagKey != "" {
			if tagTime == nil {
				span.Tags[tagKey] = tagValue
			} else {
				span.Logs = append(span.Logs, opentracing.LogData{
					Timestamp: *tagTime,
					Event:     tagKey,
					Payload:   tagValue,
				})
			}
		}
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	if span.SpanID != 0 {
		spans = append(spans, span)
	}
	return spans, nil
}

func (st *Storage) SpanWithID(id uint64) (tracer.RawSpan, error) {
	tx, err := st.db.Begin()
	if err != nil {
		return tracer.RawSpan{}, err
	}
	defer tx.Rollback()
	return st.spanWithID(tx, id)
}

func (st *Storage) spanWithID(tx *sql.Tx, id uint64) (tracer.RawSpan, error) {
	const selectSpan = `
SELECT spans.id, spans.trace_id, spans.time, spans.time, spans.operation_name, tags.key, tags.value, tags.time
FROM spans
  LEFT JOIN tags
    ON spans.id = tags.span_id
WHERE id = $1
LIMIT 1`
	rows, err := tx.Query(selectSpan, int64(id))
	if err != nil {
		return tracer.RawSpan{}, err
	}
	spans, err := scanSpans(rows)
	if err != nil {
		return tracer.RawSpan{}, err
	}
	if len(spans) == 0 {
		return tracer.RawSpan{}, sql.ErrNoRows
	}
	return spans[0], nil
}

func (st *Storage) QueryTraces(q tracer.Query) ([]tracer.RawTrace, error) {
	tx, err := st.db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	var andConds []string
	var andArgs []interface{}
	var orConds []string
	var orArgs []interface{}
	if q.FinishTime.IsZero() {
		q.FinishTime = time.Now()
	}
	for _, tag := range q.AndTags {
		if tag.CheckValue {
			andConds = append(andConds, `(tags.key = ? AND tags.value = ?)`)
			andArgs = append(andArgs, tag.Key, tag.Value)
		} else {
			andConds = append(andConds, `(tags.key = ?)`)
			andArgs = append(andArgs, tag.Key)
		}
	}

	for _, tag := range q.OrTags {
		if tag.CheckValue {
			orConds = append(orConds, `(tags.key = ? AND tags.value = ?)`)
			orArgs = append(orArgs, tag.Key, tag.Value)
		} else {
			orConds = append(orConds, `(tags.key = ?)`)
			orArgs = append(orArgs, tag.Key)
		}
	}

	and := strings.Join(andConds, " AND ")
	or := strings.Join(orConds, " OR ")
	conds := []string{"true"}
	if and != "" {
		conds = append(conds, and)
	}
	if or != "" {
		conds = append(conds, or)
	}

	query := st.db.Rebind(`
SELECT spans.trace_id
FROM spans
WHERE
  EXISTS (
    SELECT 1
    FROM tags
    WHERE
      tags.trace_id = spans.trace_id AND
      ` + strings.Join(conds, " AND ") + `
  ) AND
  ? @> spans.time AND
  (? = '' OR operation_name = ?) AND
  spans.id = spans.trace_id
ORDER BY
  spans.time ASC,
  spans.trace_id
`)
	args := make([]interface{}, 0, len(andArgs)+len(orArgs))
	args = append(args, andArgs...)
	args = append(args, orArgs...)
	args = append(args, timeRange{q.StartTime, q.FinishTime})
	args = append(args, q.OperationName, q.OperationName)

	var ids []int64
	rows, err := st.db.Query(query, args...)
	if err != nil {
		return nil, err
	}
	var id int64
	for rows.Next() {
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	var traces []tracer.RawTrace
	for _, id := range ids {
		trace, err := st.traceWithID(tx, uint64(id))
		if err != nil {
			return nil, err
		}
		traces = append(traces, trace)
	}
	return traces, nil
}
<|endoftext|>


<commit_before>package drawille

//import "code.google.com/p/goncurses"
import "math"

var pixel_map = [4][2]int{
  {0x1, 0x8},
  {0x2, 0x10},
  {0x4, 0x20},
  {0x40, 0x80}}

// Braille chars start at 0x2800
var braille_char_offset = 0x2800

func getPixel(y,x int) int {
  var cy,cx int
  if y >= 0 {
    cy = y%4
  } else {
    cy = 3 + ((y+1)%4)
  }
  if x >= 0 {
    cx = x%2
  } else {
    cx = 1 + ((x+1)%2)
  }
  return pixel_map[cy][cx]
}

type Canvas struct{
  LineEnding string
  chars map[int]map[int]int
}

// Make a new canvas
func NewCanvas() Canvas {
  c := Canvas{LineEnding:"\n"}
  c.Clear()
  return c
}

func (c Canvas) MaxY() int {
  max := 0
  for k,_ := range c.chars {
    if k>max {
      max = k
    }
  }
  return max*4
}

func (c Canvas) MinY() int {
  min := 0
  for k,_ := range c.chars {
    if k<min {
      min = k
    }
  }
  return min*4
}

func (c Canvas) MaxX() int {
  max := 0
  for _,v := range c.chars {
    for k,_ := range v {
      if k>max {
        max = k
      }
    }
  }
  return max*2
}

func (c Canvas) MinX() int {
  min := 0
  for _,v := range c.chars {
    for k,_ := range v {
      if k<min {
        min = k
      }
    }
  }
  return min*2
}

// Clear all pixels
func (c *Canvas) Clear() {
  c.chars = make(map[int]map[int]int)
}

// Convert x,y to cols, rows
func (c Canvas) get_pos(x,y int) (int,int) {
  return (x/2),(y/4)
}

// Set a pixel of c
func (c *Canvas) Set(x,y int) {
  px,py := c.get_pos(x,y)
  if m := c.chars[py];m == nil {
    c.chars[py] = make(map[int]int)
  }
  val := c.chars[py][px]
  mapv := getPixel(y,x)
  c.chars[py][px] = val|mapv
}

// Unset a pixel of c
func (c *Canvas) UnSet(x,y int) {
  px,py := c.get_pos(x,y)
  x,y = int(math.Abs(float64(x))),int(math.Abs(float64(y)))
  if m := c.chars[py]; m == nil {
    c.chars[py] = make(map[int]int)
  }
  c.chars[py][px] = c.chars[py][px]&^getPixel(y,x)
}
// Toggle a point
func (c *Canvas) Toggle(x,y int) {
  px,py := c.get_pos(x,y)
  if (c.chars[py][px] & getPixel(y,x)) != 0 {
    c.UnSet(x,y)
  } else {
    c.Set(x,y)
  }
}

// Set text to the given coordinates
func (c *Canvas) SetText(x,y int, text string) {
  x,y = x/2,y/4
  if m := c.chars[y]; m == nil {
    c.chars[y] = make(map[int]int)
  }
  for i,char := range text {
    c.chars[y][x+i] = int(char)
  }
}

// Get pixel at the given coordinates
func (c Canvas) Get(x,y int) bool {
  dot_index := pixel_map[y%4][x%2]
  x,y = x/2,y/4
  char := c.chars[y][x]
  return (char & dot_index) != 0
}

// Retrieve the rows from a given view
func (c Canvas) Rows(minX,minY,maxX,maxY int) []string {
  minrow,maxrow := minY/4,(maxY)/4
  mincol,maxcol := minX/2,(maxX)/2

  ret := make([]string,0)
  for rownum := minrow;rownum < (maxrow+1);rownum=rownum+1 {
    row := ""
    for x := mincol;x<(maxcol+1);x=x+1 {
      char := c.chars[rownum][x]
      row += string(rune(char+braille_char_offset))
    }
    ret = append(ret,row)
  }
  return ret
}

// Retrieve a string representation of the frame at the given parameters
func (c Canvas) Frame(minX,minY,maxX,maxY int) string {
  var ret string
  for _,row := range c.Rows(minX,minY,maxX,maxY) {
    ret += row
    ret += c.LineEnding
  }
  return ret
}

func (c Canvas) String() string {
  return c.Frame(c.MinX(),c.MinY(),c.MaxX(),c.MaxY())
}

func (c *Canvas) DrawLine(x1,y1, x2,y2 float64) {
  xdiff := math.Abs(x1-x2)
  ydiff := math.Abs(y2-y1)

  var xdir,ydir float64
  if x1 <= x2 {
    xdir = 1
  } else {
    xdir = -1
  }
  if y1 <= y2 {
    ydir = 1
  } else {
    ydir = -1
  }

  r := math.Max(xdiff, ydiff)

  for i:=0;i<round(r)+1;i=i+1 {
    x,y := x1,y1
    if ydiff != 0 {
      y += (float64(i)*ydiff)/(r*ydir)
    }
    if xdiff != 0 {
      x += (float64(i)*xdiff)/(r*xdir)
    }
    c.Toggle(round(x),round(y))
  }
}

func (c *Canvas) DrawPolygon(center_x,center_y,sides,radius float64) {
  degree := 360/sides
  for n:=0;n<int(sides);n=n+1 {
    a := float64(n)*degree
    b := float64(n+1)*degree

    x1 := (center_x+(math.Cos(radians(a))*(radius/2+1)))
    y1 := (center_y+(math.Sin(radians(a))*(radius/2+1)))
    x2 := (center_x+(math.Cos(radians(b))*(radius/2+1)))
    y2 := (center_y+(math.Sin(radians(b))*(radius/2+1)))

    c.DrawLine(x1,y1,x2,y2)
  }
}

func radians(d float64) float64 {
  return d*(math.Pi/180)
}

func round(x float64) int {
  return int(x+0.5)
}
<commit_msg>Fixes a bug where non braille characters are incorrectly represented when converting the int map to a string<commit_after>package drawille

//import "code.google.com/p/goncurses"
import "math"

var pixel_map = [4][2]int{
	{0x1, 0x8},
	{0x2, 0x10},
	{0x4, 0x20},
	{0x40, 0x80}}

// Braille chars start at 0x2800
var braille_char_offset = 0x2800

func getPixel(y, x int) int {
	var cy, cx int
	if y >= 0 {
		cy = y % 4
	} else {
		cy = 3 + ((y + 1) % 4)
	}
	if x >= 0 {
		cx = x % 2
	} else {
		cx = 1 + ((x + 1) % 2)
	}
	return pixel_map[cy][cx]
}

type Canvas struct {
	LineEnding string
	chars      map[int]map[int]int
}

// Make a new canvas
func NewCanvas() Canvas {
	c := Canvas{LineEnding: "\n"}
	c.Clear()
	return c
}

func (c Canvas) MaxY() int {
	max := 0
	for k, _ := range c.chars {
		if k > max {
			max = k
		}
	}
	return max * 4
}

func (c Canvas) MinY() int {
	min := 0
	for k, _ := range c.chars {
		if k < min {
			min = k
		}
	}
	return min * 4
}

func (c Canvas) MaxX() int {
	max := 0
	for _, v := range c.chars {
		for k, _ := range v {
			if k > max {
				max = k
			}
		}
	}
	return max * 2
}

func (c Canvas) MinX() int {
	min := 0
	for _, v := range c.chars {
		for k, _ := range v {
			if k < min {
				min = k
			}
		}
	}
	return min * 2
}

// Clear all pixels
func (c *Canvas) Clear() {
	c.chars = make(map[int]map[int]int)
}

// Convert x,y to cols, rows
func (c Canvas) get_pos(x, y int) (int, int) {
	return (x / 2), (y / 4)
}

// Set a pixel of c
func (c *Canvas) Set(x, y int) {
	px, py := c.get_pos(x, y)
	if m := c.chars[py]; m == nil {
		c.chars[py] = make(map[int]int)
	}
	val := c.chars[py][px]
	mapv := getPixel(y, x)
	c.chars[py][px] = val | mapv
}

// Unset a pixel of c
func (c *Canvas) UnSet(x, y int) {
	px, py := c.get_pos(x, y)
	x, y = int(math.Abs(float64(x))), int(math.Abs(float64(y)))
	if m := c.chars[py]; m == nil {
		c.chars[py] = make(map[int]int)
	}
	c.chars[py][px] = c.chars[py][px] &^ getPixel(y, x)
}

// Toggle a point
func (c *Canvas) Toggle(x, y int) {
	px, py := c.get_pos(x, y)
	if (c.chars[py][px] & getPixel(y, x)) != 0 {
		c.UnSet(x, y)
	} else {
		c.Set(x, y)
	}
}

// Set text to the given coordinates
func (c *Canvas) SetText(x, y int, text string) {
	x, y = x/2, y/4
	if m := c.chars[y]; m == nil {
		c.chars[y] = make(map[int]int)
	}
	for i, char := range text {
		c.chars[y][x+i] = int(char) - braille_char_offset
	}
}

// Get pixel at the given coordinates
func (c Canvas) Get(x, y int) bool {
	dot_index := pixel_map[y%4][x%2]
	x, y = x/2, y/4
	char := c.chars[y][x]
	return (char & dot_index) != 0
}

// Retrieve the rows from a given view
func (c Canvas) Rows(minX, minY, maxX, maxY int) []string {
	minrow, maxrow := minY/4, (maxY)/4
	mincol, maxcol := minX/2, (maxX)/2

	ret := make([]string, 0)
	for rownum := minrow; rownum < (maxrow + 1); rownum = rownum + 1 {
		row := ""
		for x := mincol; x < (maxcol + 1); x = x + 1 {
			char := c.chars[rownum][x]
			row += string(rune(char + braille_char_offset))
		}
		ret = append(ret, row)
	}
	return ret
}

// Retrieve a string representation of the frame at the given parameters
func (c Canvas) Frame(minX, minY, maxX, maxY int) string {
	var ret string
	for _, row := range c.Rows(minX, minY, maxX, maxY) {
		ret += row
		ret += c.LineEnding
	}
	return ret
}

func (c Canvas) String() string {
	return c.Frame(c.MinX(), c.MinY(), c.MaxX(), c.MaxY())
}

func (c *Canvas) DrawLine(x1, y1, x2, y2 float64) {
	xdiff := math.Abs(x1 - x2)
	ydiff := math.Abs(y2 - y1)

	var xdir, ydir float64
	if x1 <= x2 {
		xdir = 1
	} else {
		xdir = -1
	}
	if y1 <= y2 {
		ydir = 1
	} else {
		ydir = -1
	}

	r := math.Max(xdiff, ydiff)

	for i := 0; i < round(r)+1; i = i + 1 {
		x, y := x1, y1
		if ydiff != 0 {
			y += (float64(i) * ydiff) / (r * ydir)
		}
		if xdiff != 0 {
			x += (float64(i) * xdiff) / (r * xdir)
		}
		c.Toggle(round(x), round(y))
	}
}

func (c *Canvas) DrawPolygon(center_x, center_y, sides, radius float64) {
	degree := 360 / sides
	for n := 0; n < int(sides); n = n + 1 {
		a := float64(n) * degree
		b := float64(n+1) * degree

		x1 := (center_x + (math.Cos(radians(a)) * (radius/2 + 1)))
		y1 := (center_y + (math.Sin(radians(a)) * (radius/2 + 1)))
		x2 := (center_x + (math.Cos(radians(b)) * (radius/2 + 1)))
		y2 := (center_y + (math.Sin(radians(b)) * (radius/2 + 1)))

		c.DrawLine(x1, y1, x2, y2)
	}
}

func radians(d float64) float64 {
	return d * (math.Pi / 180)
}

func round(x float64) int {
	return int(x + 0.5)
}
<|endoftext|>


<commit_before>package ds_store

import (
	"bytes"
	"encoding/binary"
	"errors"
	"reflect"
	"unicode/utf16"
	"unicode/utf8"
	"unsafe"
)

type Block struct {
	Allocator *Allocator
	Offset    uint32
	Size      uint32
	Data      []byte
	Pos       uint32
}

type Allocator struct {
	Data     []byte
	Pos      uint32
	Root     *Block
	Offsets  []uint32
	Toc      map[string]uint32
	FreeList map[uint32][]uint32
}

func NewBlock(a *Allocator, pos uint32, size uint32) (block *Block, err error) {
	block = &Block{Size: size, Allocator: a, Data: a.Data[pos+0x4 : pos+0x4+size]}
	return block, nil
}

func (block *Block) readUint32() (value uint32, err error) {
	if block.Size-block.Pos < 4 {
		return 0, errors.New("Not enough bytes to read")
	}
	data := bytes.NewBuffer(block.Data)
	data.Next(int(block.Pos))
	binary.Read(data, binary.BigEndian, &value)
	block.Pos += 4
	return value, nil
}

func (block *Block) readByte() (value byte, err error) {
	if block.Size-block.Pos < 1 {
		return 0, errors.New("Not enough bytes to read")
	}
	data := bytes.NewBuffer(block.Data)
	data.Next(int(block.Pos))
	binary.Read(data, binary.BigEndian, &value)
	block.Pos += 1
	return value, nil
}

func (block *Block) readBuf(length int) (buf []byte, err error) {
	if int(block.Size)-int(block.Pos) < length {
		return nil, errors.New("Not enough bytes to read")
	}
	data := bytes.NewBuffer(block.Data)
	data.Next(int(block.Pos))
	buf = make([]byte, length)
	binary.Read(data, binary.BigEndian, &buf)
	block.Pos += uint32(length)
	return buf, nil
}

func (block *Block) readFileName() (name string, err error) {
	length, err := block.readUint32()
	if err != nil {
		return "", err
	}
	buf, err := block.readBuf(int(2 * length))
	if err != nil {
		return "", err
	}

	/*
		sid, err := block.readUint32()
		if err != nil {
			return "", err
		}
	*/
	block.skip(4)

	stype, err := block.readBuf(4)
	if err != nil {
		return "", err
	}
	t := string(stype)
	switch {
	case t == "bool":
		block.skip(1)
		break
	case t == "type" || t == "long" || t == "shor":
		block.skip(4)
		break
	case t == "comp" || t == "dutc":
		block.skip(8)
		break
	case t == "blob":
		blen, err := block.readUint32()
		if err != nil {
			return "", err
		}
		block.skip(blen)
		break
	case t == "ustr":
		blen, err := block.readUint32()
		if err != nil {
			return "", err
		}
		block.skip(2 * blen)
	default:
		panic("Unknown file format.")
	}

	name = utf16be2utf8(buf)
	return name, nil
}

func (block *Block) skip(i uint32) {
	block.Pos += i
}

func NewAllocator(data []byte) (a *Allocator, err error) {
	a = &Allocator{Data: data} //bytes.NewBuffer(data)}
	a.Toc = make(map[string]uint32)
	a.FreeList = make(map[uint32][]uint32)

	offset, size, err := a.readHeader()
	if err != nil {
		return nil, err
	}

	a.Root, _ = NewBlock(a, offset, size)

	err = a.readOffsets()
	if err != nil {
		return nil, err
	}

	err = a.readToc()
	if err != nil {
		return nil, err
	}

	err = a.readFreeList()
	if err != nil {
		return nil, err
	}

	return a, err
}

func (a *Allocator) GetBlock(bid uint32) (block *Block, err error) {
	addr := a.Offsets[bid]
	offset := int(addr) & ^0x1f
	size := 1 << (uint(addr) & 0x1f)

	block, err = NewBlock(a, uint32(offset), uint32(size)) ///+4??
	if err != nil {
		return nil, errors.New("Cannot create/read block")
	}
	return block, nil
}

func (a *Allocator) TraverseFromRootNode() (filenames []string, err error) {
	rootBlk, err := a.GetBlock(a.Toc["DSDB"])
	if err != nil {
		return nil, err
	}
	rootNode, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}
	/*height, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}
	recordsCount, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}
	nodesCount, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}
	blksize, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}*/
	rootBlk.skip(4 * 4)

	return a.Traverse(rootNode)
}

func (a *Allocator) Traverse(bid uint32) (filenames []string, err error) {
	node, err := a.GetBlock(bid)
	if err != nil {
		return nil, err
	}
	nextPtr, err := node.readUint32()
	if err != nil {
		return nil, err
	}
	count, err := node.readUint32()
	if err != nil {
		return nil, err
	}
	if nextPtr > 0 {
		//This may be broken
		for i := 0; i < int(count); i++ {
			next, err := node.readUint32()
			if err != nil {
				return nil, err
			}
			files, err := a.Traverse(next)
			if err != nil {
				return nil, err
			}
			for _, f := range files {
				filenames = append(filenames, f)
			}
			f, err := node.readFileName()
			if err != nil {
				return nil, err
			}
			filenames = append(filenames, f)
		}
		files, err := a.Traverse(nextPtr)
		if err != nil {
			return nil, err
		}
		for _, f := range files {
			filenames = append(filenames, f)
		}
	} else {
		for i := 0; i < int(count); i++ {
			f, err := node.readFileName()
			if err != nil {
				return nil, err
			}
			filenames = append(filenames, f)
		}
	}

	return filenames, nil
}

func (a *Allocator) readFreeList() error {
	for i := 0; i < 32; i++ {
		blkcount, err := a.Root.readUint32()
		if err != nil {
			return err
		}
		if blkcount == 0 {
			continue
		}
		a.FreeList[uint32(i)] = make([]uint32, 0)
		for k := 0; k < int(blkcount); k++ {
			val, err := a.Root.readUint32()
			if err != nil {
				return err
			}
			if val == 0 {
				continue
			}
			a.FreeList[uint32(i)] = append(a.FreeList[uint32(i)], val)
		}
	}
	return nil
}

func (a *Allocator) readToc() error {
	toccount, err := a.Root.readUint32()
	if err != nil {
		return err
	}
	for i := toccount; i > 0; i-- {
		tlen, err := a.Root.readByte()
		if err != nil {
			return err
		}
		name, err := a.Root.readBuf(int(tlen))
		if err != nil {
			return err
		}
		value, err := a.Root.readUint32()
		if err != nil {
			return err
		}
		a.Toc[string(name)] = value
	}
	return nil
}

func (a *Allocator) readOffsets() error {
	count, err := a.Root.readUint32()
	if err != nil {
		return err
	}
	a.Root.skip(4)

	for offcount := int(count); offcount > 0; offcount -= 256 {
		for i := 0; i < 256; i++ {
			val, err := a.Root.readUint32()
			if err != nil {
				return err
			}
			if val == 0 {
				continue
			}
			a.Offsets = append(a.Offsets, val)
		}
	}
	return nil
}

func (a *Allocator) readHeader() (offset uint32, size uint32, err error) {
	data := bytes.NewBuffer(a.Data)
	if data.Len() < 32 {
		return offset, size, errors.New("Header not long enough")
	}
	var magic1, magic, offset2 uint32

	binary.Read(data, binary.BigEndian, &magic1)
	if magic1 != 1 {
		return offset, size, errors.New("Wrong magic bytes")
	}
	a.Pos += 4

	binary.Read(data, binary.BigEndian, &magic)
	if magic != 0x42756431 {
		return offset, size, errors.New("Wrong magic bytes")
	}
	a.Pos += 4

	binary.Read(data, binary.BigEndian, &offset)
	a.Pos += 4
	binary.Read(data, binary.BigEndian, &size)
	a.Pos += 4
	binary.Read(data, binary.BigEndian, &offset2)
	if offset != offset2 {
		return offset, size, errors.New("Offets do not match")
	}
	a.Pos += 4

	return offset, size, nil
}

func utf16be2utf8(utf16be []byte) string {
	//Taken from http://play.golang.org/p/xtG1e9iqA1
	n := len(utf16be)
	// Convert to []uint16
	// hop through unsafe to skip any actual allocation/copying
	header := *(*reflect.SliceHeader)(unsafe.Pointer(&utf16be))
	header.Len /= 2
	shorts := *(*[]uint16)(unsafe.Pointer(&header))
	// shorts may need byte-swapping
	for i := 0; i < n; i += 2 {
		shorts[i/2] = (uint16(utf16be[i]) << 8) | uint16(utf16be[i+1])
	}

	// Convert to []byte
	count := 0
	for i := 0; i < len(shorts); i++ {
		r := rune(shorts[i])
		if utf16.IsSurrogate(r) {
			i++
			r = utf16.DecodeRune(r, rune(shorts[i]))
		}
		count += utf8.RuneLen(r)
	}
	buf := make([]byte, count)
	bi := 0
	for i := 0; i < len(shorts); i++ {
		r := rune(shorts[i])
		if utf16.IsSurrogate(r) {
			i++
			r = utf16.DecodeRune(r, rune(shorts[i]))
		}
		bi += utf8.EncodeRune(buf[bi:], r)
	}
	return string(buf)
}
<commit_msg>Fixed a minor bug<commit_after>package ds_store

import (
	"bytes"
	"encoding/binary"
	"errors"
	"reflect"
	"unicode/utf16"
	"unicode/utf8"
	"unsafe"
)

type Block struct {
	Allocator *Allocator
	Offset    uint32
	Size      uint32
	Data      []byte
	Pos       uint32
}

type Allocator struct {
	Data     []byte
	Pos      uint32
	Root     *Block
	Offsets  []uint32
	Toc      map[string]uint32
	FreeList map[uint32][]uint32
}

func NewBlock(a *Allocator, pos uint32, size uint32) (block *Block, err error) {
	block = &Block{Size: size, Allocator: a, Data: a.Data[pos+0x4 : pos+0x4+size]}
	return block, nil
}

func (block *Block) readUint32() (value uint32, err error) {
	if block.Size-block.Pos < 4 {
		return 0, errors.New("Not enough bytes to read")
	}
	data := bytes.NewBuffer(block.Data)
	data.Next(int(block.Pos))
	binary.Read(data, binary.BigEndian, &value)
	block.Pos += 4
	return value, nil
}

func (block *Block) readByte() (value byte, err error) {
	if block.Size-block.Pos < 1 {
		return 0, errors.New("Not enough bytes to read")
	}
	data := bytes.NewBuffer(block.Data)
	data.Next(int(block.Pos))
	binary.Read(data, binary.BigEndian, &value)
	block.Pos += 1
	return value, nil
}

func (block *Block) readBuf(length int) (buf []byte, err error) {
	if int(block.Size)-int(block.Pos) < length {
		return nil, errors.New("Not enough bytes to read")
	}
	data := bytes.NewBuffer(block.Data)
	data.Next(int(block.Pos))
	buf = make([]byte, length)
	binary.Read(data, binary.BigEndian, &buf)
	block.Pos += uint32(length)
	return buf, nil
}

func (block *Block) readFileName() (name string, err error) {
	length, err := block.readUint32()
	if err != nil {
		return "", err
	}
	buf, err := block.readBuf(int(2 * length))
	if err != nil {
		return "", err
	}

	/*
		sid, err := block.readUint32()
		if err != nil {
			return "", err
		}
	*/
	block.skip(4)

	stype, err := block.readBuf(4)
	if err != nil {
		return "", err
	}
	t := string(stype)
	switch {
	case t == "bool":
		block.skip(1)
		break
	case t == "type" || t == "long" || t == "shor":
		block.skip(4)
		break
	case t == "comp" || t == "dutc":
		block.skip(8)
		break
	case t == "blob":
		blen, err := block.readUint32()
		if err != nil {
			return "", err
		}
		block.skip(blen)
		break
	case t == "ustr":
		blen, err := block.readUint32()
		if err != nil {
			return "", err
		}
		block.skip(2 * blen)
	default:
		panic("Unknown file format.")
	}

	name = utf16be2utf8(buf)
	return name, nil
}

func (block *Block) skip(i uint32) {
	block.Pos += i
}

func NewAllocator(data []byte) (a *Allocator, err error) {
	a = &Allocator{Data: data} //bytes.NewBuffer(data)}
	a.Toc = make(map[string]uint32)
	a.FreeList = make(map[uint32][]uint32)

	offset, size, err := a.readHeader()
	if err != nil {
		return nil, err
	}

	a.Root, _ = NewBlock(a, offset, size)

	err = a.readOffsets()
	if err != nil {
		return nil, err
	}

	err = a.readToc()
	if err != nil {
		return nil, err
	}

	err = a.readFreeList()
	if err != nil {
		return nil, err
	}

	return a, err
}

func (a *Allocator) GetBlock(bid uint32) (block *Block, err error) {
	if !contains(a.Offsets, bid) {
		return nil, errors.New("Cannot find key in Offset-Table")
	}
	addr := a.Offsets[bid]

	offset := int(addr) & ^0x1f
	size := 1 << (uint(addr) & 0x1f)

	block, err = NewBlock(a, uint32(offset), uint32(size)) ///+4??
	if err != nil {
		return nil, errors.New("Cannot create/read block")
	}
	return block, nil
}

func (a *Allocator) TraverseFromRootNode() (filenames []string, err error) {
	_ = "breakpoint"
	rootBlk, err := a.GetBlock(a.Toc["DSDB"])
	if err != nil {
		return nil, err
	}
	rootNode, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}
	/*height, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}
	recordsCount, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}
	nodesCount, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}
	blksize, err := rootBlk.readUint32()
	if err != nil {
		return nil, err
	}*/
	rootBlk.skip(4 * 4)

	return a.Traverse(rootNode)
}

func (a *Allocator) Traverse(bid uint32) (filenames []string, err error) {
	node, err := a.GetBlock(bid)
	if err != nil {
		return nil, err
	}
	nextPtr, err := node.readUint32()
	if err != nil {
		return nil, err
	}
	count, err := node.readUint32()
	if err != nil {
		return nil, err
	}
	if nextPtr > 0 {
		//This may be broken
		for i := 0; i < int(count); i++ {
			next, err := node.readUint32()
			if err != nil {
				return nil, err
			}
			files, err := a.Traverse(next)
			if err != nil {
				return nil, err
			}
			for _, f := range files {
				filenames = append(filenames, f)
			}
			f, err := node.readFileName()
			if err != nil {
				return nil, err
			}
			filenames = append(filenames, f)
		}
		files, err := a.Traverse(nextPtr)
		if err != nil {
			return nil, err
		}
		for _, f := range files {
			filenames = append(filenames, f)
		}
	} else {
		for i := 0; i < int(count); i++ {
			f, err := node.readFileName()
			if err != nil {
				return nil, err
			}
			filenames = append(filenames, f)
		}
	}

	return filenames, nil
}

func (a *Allocator) readFreeList() error {
	for i := 0; i < 32; i++ {
		blkcount, err := a.Root.readUint32()
		if err != nil {
			return err
		}
		if blkcount == 0 {
			continue
		}
		a.FreeList[uint32(i)] = make([]uint32, 0)
		for k := 0; k < int(blkcount); k++ {
			val, err := a.Root.readUint32()
			if err != nil {
				return err
			}
			if val == 0 {
				continue
			}
			a.FreeList[uint32(i)] = append(a.FreeList[uint32(i)], val)
		}
	}
	return nil
}

func (a *Allocator) readToc() error {
	toccount, err := a.Root.readUint32()
	if err != nil {
		return err
	}
	for i := toccount; i > 0; i-- {
		tlen, err := a.Root.readByte()
		if err != nil {
			return err
		}
		name, err := a.Root.readBuf(int(tlen))
		if err != nil {
			return err
		}
		value, err := a.Root.readUint32()
		if err != nil {
			return err
		}
		a.Toc[string(name)] = value
	}
	return nil
}

func (a *Allocator) readOffsets() error {
	count, err := a.Root.readUint32()
	if err != nil {
		return err
	}
	a.Root.skip(4)

	for offcount := int(count); offcount > 0; offcount -= 256 {
		for i := 0; i < 256; i++ {
			val, err := a.Root.readUint32()
			if err != nil {
				return err
			}
			if val == 0 {
				continue
			}
			a.Offsets = append(a.Offsets, val)
		}
	}
	return nil
}

func (a *Allocator) readHeader() (offset uint32, size uint32, err error) {
	data := bytes.NewBuffer(a.Data)
	if data.Len() < 32 {
		return offset, size, errors.New("Header not long enough")
	}
	var magic1, magic, offset2 uint32

	binary.Read(data, binary.BigEndian, &magic1)
	if magic1 != 1 {
		return offset, size, errors.New("Wrong magic bytes")
	}
	a.Pos += 4

	binary.Read(data, binary.BigEndian, &magic)
	if magic != 0x42756431 {
		return offset, size, errors.New("Wrong magic bytes")
	}
	a.Pos += 4

	binary.Read(data, binary.BigEndian, &offset)
	a.Pos += 4
	binary.Read(data, binary.BigEndian, &size)
	a.Pos += 4
	binary.Read(data, binary.BigEndian, &offset2)
	if offset != offset2 {
		return offset, size, errors.New("Offets do not match")
	}
	a.Pos += 4

	return offset, size, nil
}

func utf16be2utf8(utf16be []byte) string {
	//Taken from http://play.golang.org/p/xtG1e9iqA1
	n := len(utf16be)
	// Convert to []uint16
	// hop through unsafe to skip any actual allocation/copying
	header := *(*reflect.SliceHeader)(unsafe.Pointer(&utf16be))
	header.Len /= 2
	shorts := *(*[]uint16)(unsafe.Pointer(&header))
	// shorts may need byte-swapping
	for i := 0; i < n; i += 2 {
		shorts[i/2] = (uint16(utf16be[i]) << 8) | uint16(utf16be[i+1])
	}

	// Convert to []byte
	count := 0
	for i := 0; i < len(shorts); i++ {
		r := rune(shorts[i])
		if utf16.IsSurrogate(r) {
			i++
			r = utf16.DecodeRune(r, rune(shorts[i]))
		}
		count += utf8.RuneLen(r)
	}
	buf := make([]byte, count)
	bi := 0
	for i := 0; i < len(shorts); i++ {
		r := rune(shorts[i])
		if utf16.IsSurrogate(r) {
			i++
			r = utf16.DecodeRune(r, rune(shorts[i]))
		}
		bi += utf8.EncodeRune(buf[bi:], r)
	}
	return string(buf)
}

//taken from http://stackoverflow.com/questions/10485743/contains-method-for-a-slice
func contains(s []uint32, e uint32) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>// data uri scheme encoder
package main

import (
	"encoding/base64"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

const DATA_URI_SCHEME_TEMPLATE = "data:image/%s;base64,"
const MAX_NOT_FLAG_COUNT = 1

// encoding to base64 data from image file
func base64encode(file os.File) string {
	fi, _ := file.Stat()
	size := fi.Size()

	data := make([]byte, size)
	file.Read(data)

	return base64.StdEncoding.EncodeToString(data)
}

// make data uri scheme string from file extention
func dataUriScheme(file os.File) string {
	fi, _ := file.Stat()
	ext := strings.Trim(filepath.Ext(fi.Name()), ".")
	return fmt.Sprintf(DATA_URI_SCHEME_TEMPLATE, ext)
}

func encode(filePath string, isPlain bool) {
	scheme := ""

	file, err := os.OpenFile(filePath, os.O_RDONLY, 0)

	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	defer func() {
		file.Close()
	}()

	if !isPlain {
		scheme = dataUriScheme(*file)
	}

	fmt.Printf("%s%s\n", scheme, base64encode(*file))
}

func main() {
	var isPlainFormat bool

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Usage of %s:
  %s [OPTIONS] FILE

Options
  -h: Show this message
`, os.Args[0], os.Args[0])
		flag.PrintDefaults()
	}
	flag.BoolVar(&isPlainFormat, "p", false, "output with plain format")
	flag.Parse()

	if flag.NArg() != MAX_NOT_FLAG_COUNT {
		flag.Usage()
		os.Exit(1)
	}

	encode(flag.Arg(0), isPlainFormat)
}
<commit_msg>Add without line break option<commit_after>// data uri scheme encoder
package main

import (
	"encoding/base64"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

const DATA_URI_SCHEME_TEMPLATE = "data:image/%s;base64,"
const MAX_NOT_FLAG_COUNT = 1

// encoding to base64 data from image file
func base64encode(file os.File) string {
	fi, _ := file.Stat()
	size := fi.Size()

	data := make([]byte, size)
	file.Read(data)

	return base64.StdEncoding.EncodeToString(data)
}

// make data uri scheme string from file extention
func dataUriScheme(file os.File) string {
	fi, _ := file.Stat()
	ext := strings.Trim(filepath.Ext(fi.Name()), ".")
	return fmt.Sprintf(DATA_URI_SCHEME_TEMPLATE, ext)
}

func encode(filePath string, isPlain bool, noRet bool) {
	scheme := ""
	ret := "\n"

	file, err := os.OpenFile(filePath, os.O_RDONLY, 0)

	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	defer func() {
		file.Close()
	}()

	if !isPlain {
		scheme = dataUriScheme(*file)
	}

	if noRet {
		ret = ""
	}

	fmt.Printf("%s%s%s", scheme, base64encode(*file), ret)
}

func main() {
	var isPlainFormat bool
	var noRet bool

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Usage of %s:
  %s [OPTIONS] FILE

Options
  -h: Show this message
`, os.Args[0], os.Args[0])
		flag.PrintDefaults()
	}
	flag.BoolVar(&isPlainFormat, "p", false, "output with plain format")
	flag.BoolVar(&noRet, "n", false, "without line break")
	flag.Parse()

	if flag.NArg() != MAX_NOT_FLAG_COUNT {
		flag.Usage()
		os.Exit(1)
	}

	encode(flag.Arg(0), isPlainFormat, noRet)
}
<|endoftext|>