<commit_before>// (c) Copyright 2015-2017 JONNALAGADDA Srinivas
//
// 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 flow

import (
	"database/sql"
	"errors"
	"fmt"
	"math"
	"time"
)

// EventStatus enumerates the query parameter values for filtering by
// event state.
type EventStatus uint8

const (
	EventStatusAll EventStatus = iota
	EventStatusApplied
	EventStatusPending
)

// DocEventID is the type of unique document event identifiers.
type DocEventID int64

// DocEvent represents a user action performed on a document in the
// system.
//
// Together with documents and nodes, events are central to the
// workflow engine in `flow`.  Events cause documents to transition
// from one state to another, usually in response to user actions.  It
// is possible for system events to cause state transitions, as well.
type DocEvent struct {
	id     DocEventID  // Unique ID of this event
	docID  DocumentID  // Document to which this event is to be applied
	state  DocState    // Current state of the document
	user   UserID      // User who caused this action
	action DocAction   // Action performed by the user
	text   string      // Comment or other content
	ctime  time.Time   // Time at which the event occurred
	status EventStatus // Status of this event
}

// ID answers the unique identifier of this document event.
func (e *DocEvent) ID() DocEventID {
	return e.id
}

// Document answers the document which this event transformed.
func (e *DocEvent) Document() DocumentID {
	return e.docID
}

// State answers the state of the document as of the time this event
// occurred.
func (e *DocEvent) State() DocState {
	return e.state
}

// User answers the user who caused this event to occur.
func (e *DocEvent) User() UserID {
	return e.user
}

// Action answers the document action that this event represents.
func (e *DocEvent) Action() DocAction {
	return e.action
}

// Text answers the comment or other content enclosed in this event.
func (e *DocEvent) Text() string {
	return e.text
}

// Time answers the time at which this event occurred.
func (e *DocEvent) Time() time.Time {
	return e.ctime
}

// Status answers the status of this event.
func (e *DocEvent) Status() (EventStatus, error) {
	var dstatus string
	row := db.QueryRow("SELECT status FROM wf_docevents WHERE id = ?", e.id)
	err := row.Scan(&dstatus)
	if err != nil {
		return 0, err
	}
	switch dstatus {
	case "A":
		e.status = EventStatusApplied

	case "P":
		e.status = EventStatusPending

	default:
		return 0, fmt.Errorf("unknown event status : %s", dstatus)
	}

	return e.status, nil
}

// Unexported type, only for convenience methods.
type _DocEvents struct{}

var _docevents *_DocEvents

func init() {
	_docevents = &_DocEvents{}
}

// New creates and initialises an event that transforms the document
// that it refers to.
func (des *_DocEvents) New(otx *sql.Tx, user UserID, doc *Document, action DocAction, text string) (DocEventID, error) {
	if doc == nil || user <= 0 {
		return 0, errors.New("document should be non-nil; user ID should be a positive integer")
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	INSERT INTO wf_docevents(doc_id, docstate_id, user_id, docaction_id, data, ctime, status)
	VALUES(?, ?, ?, ?, ?, NOW(), 'P')
	`
	res, err := tx.Exec(q, doc.id, doc.state, user, action.id, text)
	if err != nil {
		return 0, err
	}
	var id int64
	id, err = res.LastInsertId()
	if err != nil {
		return 0, err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return DocEventID(id), nil
}

// List answers a subset of document events, based on the input
// specification.
//
// `status` should be one of `all`, `applied` and `pending`.
//
// Result set begins with ID >= `offset`, and has not more than
// `limit` elements.  A value of `0` for `offset` fetches from the
// beginning, while a value of `0` for `limit` fetches until the end.
func (des *_DocEvents) List(status EventStatus, offset, limit int64) ([]*DocEvent, error) {
	if offset < 0 || limit < 0 {
		return nil, errors.New("offset and limit must be non-negative integers")
	}
	if limit == 0 {
		limit = math.MaxInt64
	}

	q := `
	SELECT *
	FROM wf_docevents
	`
	switch status {
	case EventStatusAll:
		// Intentionally left blank

	case EventStatusApplied:
		q = q + `
		WHERE status = 'A'
		`

	case EventStatusPending:
		q = q + `
		WHERE status = 'P'
		`

	default:
		return nil, fmt.Errorf("unknown event status specified in filter : %d", status)
	}
	q = q + `
	ORDER BY id
	LIMIT ? OFFSET ?
	`
	rows, err := db.Query(q, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var text sql.NullString
	var dstatus string
	ary := make([]*DocEvent, 0, 10)
	for rows.Next() {
		var elem DocEvent
		err = rows.Scan(&elem.id, &elem.docID, &elem.state, &elem.user, &elem.action, &text, &elem.ctime, &dstatus)
		if err != nil {
			return nil, err
		}
		if text.Valid {
			elem.text = text.String
		}
		switch dstatus {
		case "A":
			elem.status = EventStatusApplied

		case "P":
			elem.status = EventStatusPending

		default:
			return nil, fmt.Errorf("unknown event status : %s", dstatus)
		}
		ary = append(ary, &elem)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return ary, nil
}

// Get retrieves a document event from the database, using the given
// event ID.
func (des *_DocEvents) Get(eid DocEventID) (*DocEvent, error) {
	if eid <= 0 {
		return nil, errors.New("event ID should be a positive integer")
	}

	var text sql.NullString
	var dstatus string
	var elem DocEvent
	row := db.QueryRow("SELECT * FROM wf_docevents WHERE id = ?", eid)
	err := row.Scan(&elem.id, &elem.docID, &elem.state, &elem.user, &elem.action, &text, &elem.ctime, &dstatus)
	if err != nil {
		return nil, err
	}
	if text.Valid {
		elem.text = text.String
	}
	switch dstatus {
	case "A":
		elem.status = EventStatusApplied

	case "P":
		elem.status = EventStatusPending

	default:
		return nil, fmt.Errorf("unknown event status : %s", dstatus)
	}

	return &elem, nil
}
<commit_msg>Modify `DocEvent` creation to take appropriate IDs as parameters<commit_after>// (c) Copyright 2015-2017 JONNALAGADDA Srinivas
//
// 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 flow

import (
	"database/sql"
	"errors"
	"fmt"
	"math"
	"time"
)

// EventStatus enumerates the query parameter values for filtering by
// event state.
type EventStatus uint8

const (
	EventStatusAll EventStatus = iota
	EventStatusApplied
	EventStatusPending
)

// DocEventID is the type of unique document event identifiers.
type DocEventID int64

// DocEvent represents a user action performed on a document in the
// system.
//
// Together with documents and nodes, events are central to the
// workflow engine in `flow`.  Events cause documents to transition
// from one state to another, usually in response to user actions.  It
// is possible for system events to cause state transitions, as well.
type DocEvent struct {
	id     DocEventID  // Unique ID of this event
	docID  DocumentID  // Document to which this event is to be applied
	state  DocState    // Current state of the document
	user   UserID      // User who caused this action
	action DocAction   // Action performed by the user
	text   string      // Comment or other content
	ctime  time.Time   // Time at which the event occurred
	status EventStatus // Status of this event
}

// ID answers the unique identifier of this document event.
func (e *DocEvent) ID() DocEventID {
	return e.id
}

// Document answers the document which this event transformed.
func (e *DocEvent) Document() DocumentID {
	return e.docID
}

// State answers the state of the document as of the time this event
// occurred.
func (e *DocEvent) State() DocState {
	return e.state
}

// User answers the user who caused this event to occur.
func (e *DocEvent) User() UserID {
	return e.user
}

// Action answers the document action that this event represents.
func (e *DocEvent) Action() DocAction {
	return e.action
}

// Text answers the comment or other content enclosed in this event.
func (e *DocEvent) Text() string {
	return e.text
}

// Time answers the time at which this event occurred.
func (e *DocEvent) Time() time.Time {
	return e.ctime
}

// Status answers the status of this event.
func (e *DocEvent) Status() (EventStatus, error) {
	var dstatus string
	row := db.QueryRow("SELECT status FROM wf_docevents WHERE id = ?", e.id)
	err := row.Scan(&dstatus)
	if err != nil {
		return 0, err
	}
	switch dstatus {
	case "A":
		e.status = EventStatusApplied

	case "P":
		e.status = EventStatusPending

	default:
		return 0, fmt.Errorf("unknown event status : %s", dstatus)
	}

	return e.status, nil
}

// Unexported type, only for convenience methods.
type _DocEvents struct{}

var _docevents *_DocEvents

func init() {
	_docevents = &_DocEvents{}
}

// New creates and initialises an event that transforms the document
// that it refers to.
func (des *_DocEvents) New(otx *sql.Tx, user UserID, dtype DocTypeID, did DocumentID,
	state DocStateID, action DocActionID, text string) (DocEventID, error) {
	if user <= 0 {
		return 0, errors.New("user ID should be a positive integer")
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	INSERT INTO wf_docevents(doctype_id, doc_id, docstate_id, docaction_id, user_id, data, ctime, status)
	VALUES(?, ?, ?, ?, ?, NOW(), 'P')
	`
	res, err := tx.Exec(q, dtype, did, state, action, text)
	if err != nil {
		return 0, err
	}
	var id int64
	id, err = res.LastInsertId()
	if err != nil {
		return 0, err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return DocEventID(id), nil
}

// List answers a subset of document events, based on the input
// specification.
//
// `status` should be one of `all`, `applied` and `pending`.
//
// Result set begins with ID >= `offset`, and has not more than
// `limit` elements.  A value of `0` for `offset` fetches from the
// beginning, while a value of `0` for `limit` fetches until the end.
func (des *_DocEvents) List(status EventStatus, offset, limit int64) ([]*DocEvent, error) {
	if offset < 0 || limit < 0 {
		return nil, errors.New("offset and limit must be non-negative integers")
	}
	if limit == 0 {
		limit = math.MaxInt64
	}

	q := `
	SELECT *
	FROM wf_docevents
	`
	switch status {
	case EventStatusAll:
		// Intentionally left blank

	case EventStatusApplied:
		q = q + `
		WHERE status = 'A'
		`

	case EventStatusPending:
		q = q + `
		WHERE status = 'P'
		`

	default:
		return nil, fmt.Errorf("unknown event status specified in filter : %d", status)
	}
	q = q + `
	ORDER BY id
	LIMIT ? OFFSET ?
	`
	rows, err := db.Query(q, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var text sql.NullString
	var dstatus string
	ary := make([]*DocEvent, 0, 10)
	for rows.Next() {
		var elem DocEvent
		err = rows.Scan(&elem.id, &elem.docID, &elem.state, &elem.user, &elem.action, &text, &elem.ctime, &dstatus)
		if err != nil {
			return nil, err
		}
		if text.Valid {
			elem.text = text.String
		}
		switch dstatus {
		case "A":
			elem.status = EventStatusApplied

		case "P":
			elem.status = EventStatusPending

		default:
			return nil, fmt.Errorf("unknown event status : %s", dstatus)
		}
		ary = append(ary, &elem)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return ary, nil
}

// Get retrieves a document event from the database, using the given
// event ID.
func (des *_DocEvents) Get(eid DocEventID) (*DocEvent, error) {
	if eid <= 0 {
		return nil, errors.New("event ID should be a positive integer")
	}

	var text sql.NullString
	var dstatus string
	var elem DocEvent
	row := db.QueryRow("SELECT * FROM wf_docevents WHERE id = ?", eid)
	err := row.Scan(&elem.id, &elem.docID, &elem.state, &elem.user, &elem.action, &text, &elem.ctime, &dstatus)
	if err != nil {
		return nil, err
	}
	if text.Valid {
		elem.text = text.String
	}
	switch dstatus {
	case "A":
		elem.status = EventStatusApplied

	case "P":
		elem.status = EventStatusPending

	default:
		return nil, fmt.Errorf("unknown event status : %s", dstatus)
	}

	return &elem, nil
}
<|endoftext|>


<commit_before>// This is Free Software covered by the terms of the MIT license.
// See LICENSE file for details.
// Copyright 2017 by Intevation GmbH
package main

type nodeType int

const (
	orderedListNode nodeType = iota
	unorderedListNode
	listItemNode
	textNode
	boldNode
	italicsNode
	underlinedNode
	strikeNode
	superscriptNode
	subscriptNode
	tableNode
	tableHeaderRowNode
	tableHeaderCellNode
	tableRowNode
	tableCellNode
	heading1Node
	heading2Node
	heading3Node
	heading4Node
	heading5Node
	heading6Node
	paragraphNode
	lineBreakNode
	escapeNode
	noWikiNode
	noWikiInlineNode
	placeholderNode
	imageNode
	linkNode
	horizontalLineNode
)

type node struct {
	nodeType

	parent   *node
	children []*node
	value    interface{}
}

type image struct {
	src string
	alt string
}

func (n *node) nextOfType(c *node) *node {
	if n == nil {
		return nil
	}
	var found bool
	for _, x := range n.children {
		if found {
			if x.nodeType == c.nodeType {
				return x
			}
		} else {
			found = x == c
		}
	}
	return nil
}

func (n *node) numColumns() int {
	var max int
	var descent func(*node)
	descent = func(current *node) {
		if current == nil {
			return
		}
		if current.nodeType == tableRowNode ||
			current.nodeType == tableHeaderRowNode {
			var count int
			for _, c := range current.children {
				if c.nodeType == tableCellNode ||
					c.nodeType == tableHeaderCellNode {
					count++
				}
			}
			if count > max {
				max = count
			}
			return
		}
		for _, c := range current.children {
			descent(c)
		}
	}
	descent(n)
	return max
}

type document struct {
	root *node
}

type visitor struct {
	enter func(*node) error
	leave func(*node) error
}

func (d *document) traverse(visitors map[nodeType]*visitor) error {

	if d == nil {
		return nil
	}

	var decent func(*node) error
	decent = func(n *node) error {
		if n == nil {
			return nil
		}
		v := visitors[n.nodeType]
		if v != nil && v.enter != nil {
			if err := v.enter(n); err != nil {
				return err
			}
		}

		for _, child := range n.children {
			if err := decent(child); err != nil {
				return err
			}
		}

		if v != nil && v.leave != nil {
			if err := v.leave(n); err != nil {
				return err
			}
		}
		return nil
	}

	return decent(d.root)
}
<commit_msg>Spell descend correctly.<commit_after>// This is Free Software covered by the terms of the MIT license.
// See LICENSE file for details.
// Copyright 2017 by Intevation GmbH
package main

type nodeType int

const (
	orderedListNode nodeType = iota
	unorderedListNode
	listItemNode
	textNode
	boldNode
	italicsNode
	underlinedNode
	strikeNode
	superscriptNode
	subscriptNode
	tableNode
	tableHeaderRowNode
	tableHeaderCellNode
	tableRowNode
	tableCellNode
	heading1Node
	heading2Node
	heading3Node
	heading4Node
	heading5Node
	heading6Node
	paragraphNode
	lineBreakNode
	escapeNode
	noWikiNode
	noWikiInlineNode
	placeholderNode
	imageNode
	linkNode
	horizontalLineNode
)

type node struct {
	nodeType

	parent   *node
	children []*node
	value    interface{}
}

type image struct {
	src string
	alt string
}

func (n *node) nextOfType(c *node) *node {
	if n == nil {
		return nil
	}
	var found bool
	for _, x := range n.children {
		if found {
			if x.nodeType == c.nodeType {
				return x
			}
		} else {
			found = x == c
		}
	}
	return nil
}

func (n *node) numColumns() int {
	var max int
	var descend func(*node)
	descend = func(current *node) {
		if current == nil {
			return
		}
		if current.nodeType == tableRowNode ||
			current.nodeType == tableHeaderRowNode {
			var count int
			for _, c := range current.children {
				if c.nodeType == tableCellNode ||
					c.nodeType == tableHeaderCellNode {
					count++
				}
			}
			if count > max {
				max = count
			}
			return
		}
		for _, c := range current.children {
			descend(c)
		}
	}
	descend(n)
	return max
}

type document struct {
	root *node
}

type visitor struct {
	enter func(*node) error
	leave func(*node) error
}

func (d *document) traverse(visitors map[nodeType]*visitor) error {

	if d == nil {
		return nil
	}

	var descend func(*node) error
	descend = func(n *node) error {
		if n == nil {
			return nil
		}
		v := visitors[n.nodeType]
		if v != nil && v.enter != nil {
			if err := v.enter(n); err != nil {
				return err
			}
		}

		for _, child := range n.children {
			if err := descend(child); err != nil {
				return err
			}
		}

		if v != nil && v.leave != nil {
			if err := v.leave(n); err != nil {
				return err
			}
		}
		return nil
	}

	return descend(d.root)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"context"
	"encoding/base64"
	"fmt"
	"github.com/hanjm/log"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strings"
	"time"
)

type Task interface {
	Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) error
	IsCompleted() bool
	FileName() string
	ContentLength() int64
}

func NewDownloadTask(sourceURL string) (Task, error) {
	switch {
	case strings.HasPrefix(sourceURL, "http"):
		return NewHTTPTask(sourceURL), nil
	case strings.HasPrefix(sourceURL, "magnet:?xt=urn:btih:") || IsBase64String(sourceURL):
		return NewMagnetTask(sourceURL), nil
	default:
		return nil, fmt.Errorf("sourceURL expect http or magnet, not %s", sourceURL)
	}
}

const (
	DownloadTaskTypeHTTP = iota
	DownloadTaskTypeMagnet
)

type TaskInfo struct {
	SourceURL     string
	StartTime     time.Time
	FileName      string
	ContentLength int64         // B 总大小
	Size          int64         // B 已下载的大小
	Duration      time.Duration // s 耗时
	Speed         int64         // B/s 速度
	IsCompleted   bool          // 是否完成
	IsError       bool          // 是否出错
	Error         string        // 错误消息
}

// download http content
type HTTPTask struct {
	TaskType int
	TaskInfo
}

func NewHTTPTask(sourceUrl string) *HTTPTask {
	return &HTTPTask{
		TaskType: DownloadTaskTypeHTTP,
		TaskInfo: TaskInfo{
			SourceURL: sourceUrl,
			FileName:  getSafeFilename(sourceUrl),
		}}
}

func (t *HTTPTask) Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) error {
	var httpClient = &http.Client{
		Timeout: limitTimeout,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 20 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 20 * time.Second,
		},
	}
	t.StartTime = time.Now()
	resp, err := httpClient.Get(t.SourceURL)
	if err != nil {
		return t.Errorf("http.Client error:%s", err)
	}
	defer resp.Body.Close()
	t.TaskInfo.ContentLength = resp.ContentLength
	contentDisposition := strings.SplitN(resp.Header.Get("Content-Disposition"), "=", 2)
	var attachmentName string
	if len(contentDisposition) > 1 {
		attachmentName = contentDisposition[1]
	}
	if t.TaskInfo.ContentLength <= 0 {
		resp.Body.Close()
		//一些资源是动态生成的,请求第一次是chunked stream,Header不带Content-Length,第二次请求就有Content-length
		resp, err = httpClient.Get(t.SourceURL)
		if err != nil {
			return t.Errorf("http.Client error:%s", err)
		}
		t.TaskInfo.ContentLength = resp.ContentLength
	}
	// if header has attach filename, update it
	if attachmentName != "" {
		attachmentName2 := getSafeFilename(attachmentName)
		if attachmentName2 != t.TaskInfo.FileName {
			t.TaskInfo.FileName = attachmentName2
		}
	}
	log.Infof("create HTTP task: length:%s source:%s filename:%s", getHumanSizeString(t.TaskInfo.ContentLength), t.SourceURL, t.TaskInfo.FileName)
	if t.TaskInfo.ContentLength > limitByteSize {
		return t.Errorf("the content length of sourceUrl is too big:%d, limit:%d", t.TaskInfo.ContentLength, limitByteSize)
	}
	// write file
	filename := downloadDir + "/" + t.TaskInfo.FileName
	os.Remove(filename)
	fp, err := os.Create(filename)
	if err != nil {
		return t.Errorf("create file error:%s", err)
	}
	defer fp.Sync()
	defer fp.Close()
	bufSize := 4096
	bodyReader := bufio.NewReaderSize(resp.Body, bufSize)
	if err != nil {
		return t.Errorf("create file error:%s", err)
	}
	buf := make([]byte, bufSize)
	size := 0
	readSize := 0
	completed := false
	for i := 0; ; i++ {
		readSize, err = bodyReader.Read(buf)
		if err != nil {
			if err == io.EOF {
				// 正常下载完
				completed = true
			} else {
				return t.Errorf("body read error:%s", err)
			}
		}
		_, err = fp.Write(buf[:readSize])
		size += readSize
		t.Size = int64(size)
		if i%1000 == 0 {
			t.Duration = time.Now().Sub(t.StartTime)
			t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
		}
		if err != nil {
			return t.Errorf("body write error:%s", err)
		}
		if completed {
			break
		}
	}
	t.TaskInfo.IsCompleted = true
	t.Size = int64(size)
	t.TaskInfo.ContentLength = t.Size
	t.Duration = time.Now().Sub(t.StartTime)
	t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
	log.Infof("complete HTTP task: length:%s source:%s filename:%s, duration:%s", getHumanSizeString(t.TaskInfo.ContentLength), t.SourceURL, t.TaskInfo.FileName, t.Duration)
	return nil
}

func (t *HTTPTask) IsCompleted() bool {
	return t.TaskInfo.IsCompleted
}

func (t *HTTPTask) FileName() string {
	return t.TaskInfo.FileName
}

func (t *HTTPTask) ContentLength() int64 {
	return t.TaskInfo.ContentLength
}

func (t *HTTPTask) Errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	_, file, line, ok := runtime.Caller(1)
	if ok {
		log.Errorf("[%s:%d]%s", file, line, err.Error())
	}
	t.TaskInfo.IsError = true
	t.TaskInfo.IsCompleted = true
	t.TaskInfo.Error = err.Error()
	return err
}

// download magnet content
type MagnetTask struct {
	TaskType int
	TaskInfo
}

func NewMagnetTask(sourceUrl string) *MagnetTask {
	return &MagnetTask{
		TaskType: DownloadTaskTypeMagnet,
		TaskInfo: TaskInfo{
			SourceURL: sourceUrl,
			FileName:  getSafeFilename(sourceUrl),
		}}
}
func (t *MagnetTask) Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) (err error) {
	if !IsAria2cRunning() {
		return t.Errorf("aria2c is not running, cannot download magnet")
	}
	aria2cRPCClient := NewAria2cRPCClient()
	var taskGID string
	// magnet? / torrent? / torrent file in downloadDir
	var isMagnetLink bool
	var torrentBase64 string
	data, err := base64.StdEncoding.DecodeString(t.SourceURL)
	if err != nil {
		if data, err = ioutil.ReadFile(t.SourceURL); err != nil {
			isMagnetLink = true
		} else {
			// try read from torrent file in downloadDir, for reDownload torrent
			isMagnetLink = false
			torrentBase64 = base64.StdEncoding.EncodeToString(data)
		}
	} else {
		isMagnetLink = false
		torrentBase64 = t.SourceURL
	}
	if isMagnetLink {
		taskGID, err = aria2cRPCClient.AddURI(torrentBase64)
		if err != nil {
			return t.Errorf("call aria2c AddURI error:%s", err)
		}
	} else {
		taskGID, err = aria2cRPCClient.AddTorrent(torrentBase64)
		if err != nil {
			return t.Errorf("call aria2c AddTorrent error:%s", err)
		}
		// save to file and change the sourceURL
		torrentFilename := fmt.Sprintf("%s/%s.torrent", downloadDir, torrentBase64[:16])
		fp, err := os.Create(torrentFilename)
		if err != nil {
			log.Warnf("os.Create file error:%s", err)
		}
		_, err = fp.Write(data)
		if err != nil {
			log.Warnf("fp.Write error:%s", err)
		}
		t.SourceURL = torrentFilename
	}
	log.Infof("create Magnet task: sourceURL:%s, taskGID:%s", t.SourceURL, taskGID)
	t.StartTime = time.Now()
	timeout := limitTimeout
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
MagnetLoop:
	complete := false
	ticker := time.NewTicker(time.Second * 5)
	for !complete {
		select {
		case <-ticker.C:
			result, err := aria2cRPCClient.TellStatus(taskGID)
			if err != nil {
				return t.Errorf("call aria2c TellStatus error:%s", err)
			}
			// update break condition
			complete = result.Completed()
			// udpate taskInfo
			t.TaskInfo.ContentLength = result.TotalLength
			t.Size = result.CompletedLength
			t.Duration = time.Now().Sub(t.StartTime)
			t.Speed = result.DownloadSpeed
			if !complete && result.CompletedLength > 0 && result.CompletedLength == result.TotalLength {
				// why aria2c wait so long time even it seems the download task is completed, may be as a seeder?
				log.Infof("force to set task status complete, status:%+v", result)
				time.Sleep(time.Second * 5)
				log.Infof("force to set task status complete, status:%+v", result)
				complete = true
			}
			// 磁力链接建立任务时无法指定文件名 获得真实文件名后需要重命名
			realFilename := strings.TrimPrefix(result.GetFilePath(), downloadDir+"/")
			pos := strings.Index(realFilename, "/")
			if pos != -1 && pos < len(realFilename)-1 {
				realFilename = realFilename[:pos]
			}
			t.TaskInfo.FileName = realFilename
			// 检查是否有继续下载磁力链接包含的其他文件
			followedBys := result.FollowedBy
			for _, followedTaskGID := range followedBys {
				taskGID = followedTaskGID
				complete = false
				log.Debugf("goto MagnetLoop: task status:%+v", result)
				goto MagnetLoop
			}
		case <-ctx.Done():
			if ctx.Err() == context.DeadlineExceeded {
				return t.Errorf("task timeout:%s", timeout)
			}
		}
	}
	err = aria2cRPCClient.RemoveDownloadResult(taskGID)
	if err != nil {
		log.Warnf("RemoveDownloadResult error:%s", err)
	}
	t.TaskInfo.IsCompleted = true
	t.Duration = time.Now().Sub(t.StartTime)
	t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
	return nil
}

func (t *MagnetTask) IsCompleted() bool {
	return t.TaskInfo.IsCompleted
}

func (t *MagnetTask) FileName() string {
	return t.TaskInfo.FileName
}

func (t *MagnetTask) ContentLength() int64 {
	return t.TaskInfo.ContentLength
}

func (t *MagnetTask) Errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	_, file, line, ok := runtime.Caller(1)
	if ok {
		log.Errorf("[%s:%d]%s", file, line, err.Error())
	}
	t.TaskInfo.IsError = true
	t.TaskInfo.IsCompleted = true
	t.TaskInfo.Error = err.Error()
	return err
}

//utils
var safeFilenameRegexp = regexp.MustCompile(`[\w\d\-.]+`)

func getSafeFilename(str string) string {
	filename := strings.Join(safeFilenameRegexp.FindAllString(str, -1), "")
	if lenOfFilename := len(filename); lenOfFilename > 50 {
		filename = filename[lenOfFilename-50 : lenOfFilename]
	}
	return fmt.Sprintf("%d-%s", time.Now().Unix(), filename)
}

func getHumanSizeString(byteSize int64) string {
	units := []string{"B", "KB", "MB", "GB", "TB", "EB"}
	index := 0
	byteSizeFloat := float64(byteSize)
	for ; byteSizeFloat > 1024; index += 1 {
		byteSizeFloat /= 1024
	}
	var unit string
	if index < len(units) {
		unit = units[index]
	} else {
		unit = "INF"
	}
	return fmt.Sprintf("%.2f %s", byteSizeFloat, unit)
}

func calculateDownloadSpeed(byteSize int64, duration time.Duration) (bytePerSecond int64) {
	if duration <= 0 {
		duration = 1
	}
	return byteSize * 1e9 / int64(duration)
}

func IsBase64String(str string) bool {
	_, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return false
	}
	return true
}
<commit_msg>fix download magnetlink url<commit_after>package main

import (
	"bufio"
	"context"
	"encoding/base64"
	"fmt"
	"github.com/hanjm/log"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strings"
	"time"
)

type Task interface {
	Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) error
	IsCompleted() bool
	FileName() string
	ContentLength() int64
}

func NewDownloadTask(sourceURL string) (Task, error) {
	switch {
	case strings.HasPrefix(sourceURL, "http"):
		return NewHTTPTask(sourceURL), nil
	case strings.HasPrefix(sourceURL, "magnet:?xt=urn:btih:") || IsBase64String(sourceURL):
		return NewMagnetTask(sourceURL), nil
	default:
		return nil, fmt.Errorf("sourceURL expect http or magnet, not %s", sourceURL)
	}
}

const (
	DownloadTaskTypeHTTP = iota
	DownloadTaskTypeMagnet
)

type TaskInfo struct {
	SourceURL     string
	StartTime     time.Time
	FileName      string
	ContentLength int64         // B 总大小
	Size          int64         // B 已下载的大小
	Duration      time.Duration // s 耗时
	Speed         int64         // B/s 速度
	IsCompleted   bool          // 是否完成
	IsError       bool          // 是否出错
	Error         string        // 错误消息
}

// download http content
type HTTPTask struct {
	TaskType int
	TaskInfo
}

func NewHTTPTask(sourceUrl string) *HTTPTask {
	return &HTTPTask{
		TaskType: DownloadTaskTypeHTTP,
		TaskInfo: TaskInfo{
			SourceURL: sourceUrl,
			FileName:  getSafeFilename(sourceUrl),
		}}
}

func (t *HTTPTask) Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) error {
	var httpClient = &http.Client{
		Timeout: limitTimeout,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 20 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 20 * time.Second,
		},
	}
	t.StartTime = time.Now()
	resp, err := httpClient.Get(t.SourceURL)
	if err != nil {
		return t.Errorf("http.Client error:%s", err)
	}
	defer resp.Body.Close()
	t.TaskInfo.ContentLength = resp.ContentLength
	contentDisposition := strings.SplitN(resp.Header.Get("Content-Disposition"), "=", 2)
	var attachmentName string
	if len(contentDisposition) > 1 {
		attachmentName = contentDisposition[1]
	}
	if t.TaskInfo.ContentLength <= 0 {
		resp.Body.Close()
		//一些资源是动态生成的,请求第一次是chunked stream,Header不带Content-Length,第二次请求就有Content-length
		resp, err = httpClient.Get(t.SourceURL)
		if err != nil {
			return t.Errorf("http.Client error:%s", err)
		}
		t.TaskInfo.ContentLength = resp.ContentLength
	}
	// if header has attach filename, update it
	if attachmentName != "" {
		attachmentName2 := getSafeFilename(attachmentName)
		if attachmentName2 != t.TaskInfo.FileName {
			t.TaskInfo.FileName = attachmentName2
		}
	}
	log.Infof("create HTTP task: length:%s source:%s filename:%s", getHumanSizeString(t.TaskInfo.ContentLength), t.SourceURL, t.TaskInfo.FileName)
	if t.TaskInfo.ContentLength > limitByteSize {
		return t.Errorf("the content length of sourceUrl is too big:%d, limit:%d", t.TaskInfo.ContentLength, limitByteSize)
	}
	// write file
	filename := downloadDir + "/" + t.TaskInfo.FileName
	os.Remove(filename)
	fp, err := os.Create(filename)
	if err != nil {
		return t.Errorf("create file error:%s", err)
	}
	defer fp.Sync()
	defer fp.Close()
	bufSize := 4096
	bodyReader := bufio.NewReaderSize(resp.Body, bufSize)
	if err != nil {
		return t.Errorf("create file error:%s", err)
	}
	buf := make([]byte, bufSize)
	size := 0
	readSize := 0
	completed := false
	for i := 0; ; i++ {
		readSize, err = bodyReader.Read(buf)
		if err != nil {
			if err == io.EOF {
				// 正常下载完
				completed = true
			} else {
				return t.Errorf("body read error:%s", err)
			}
		}
		_, err = fp.Write(buf[:readSize])
		size += readSize
		t.Size = int64(size)
		if i%1000 == 0 {
			t.Duration = time.Now().Sub(t.StartTime)
			t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
		}
		if err != nil {
			return t.Errorf("body write error:%s", err)
		}
		if completed {
			break
		}
	}
	t.TaskInfo.IsCompleted = true
	t.Size = int64(size)
	t.TaskInfo.ContentLength = t.Size
	t.Duration = time.Now().Sub(t.StartTime)
	t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
	log.Infof("complete HTTP task: length:%s source:%s filename:%s, duration:%s", getHumanSizeString(t.TaskInfo.ContentLength), t.SourceURL, t.TaskInfo.FileName, t.Duration)
	return nil
}

func (t *HTTPTask) IsCompleted() bool {
	return t.TaskInfo.IsCompleted
}

func (t *HTTPTask) FileName() string {
	return t.TaskInfo.FileName
}

func (t *HTTPTask) ContentLength() int64 {
	return t.TaskInfo.ContentLength
}

func (t *HTTPTask) Errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	_, file, line, ok := runtime.Caller(1)
	if ok {
		log.Errorf("[%s:%d]%s", file, line, err.Error())
	}
	t.TaskInfo.IsError = true
	t.TaskInfo.IsCompleted = true
	t.TaskInfo.Error = err.Error()
	return err
}

// download magnet content
type MagnetTask struct {
	TaskType int
	TaskInfo
}

func NewMagnetTask(sourceUrl string) *MagnetTask {
	return &MagnetTask{
		TaskType: DownloadTaskTypeMagnet,
		TaskInfo: TaskInfo{
			SourceURL: sourceUrl,
			FileName:  getSafeFilename(sourceUrl),
		}}
}
func (t *MagnetTask) Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) (err error) {
	if !IsAria2cRunning() {
		return t.Errorf("aria2c is not running, cannot download magnet")
	}
	aria2cRPCClient := NewAria2cRPCClient()
	var taskGID string
	// magnet? / torrent? / torrent file in downloadDir
	var isMagnetLink bool
	var torrentBase64 string
	data, err := base64.StdEncoding.DecodeString(t.SourceURL)
	if err != nil {
		if data, err = ioutil.ReadFile(t.SourceURL); err != nil {
			isMagnetLink = true
		} else {
			// try read from torrent file in downloadDir, for reDownload torrent
			isMagnetLink = false
			torrentBase64 = base64.StdEncoding.EncodeToString(data)
		}
	} else {
		isMagnetLink = false
		torrentBase64 = t.SourceURL
	}
	if isMagnetLink {
		taskGID, err = aria2cRPCClient.AddURI(t.SourceURL)
		if err != nil {
			return t.Errorf("call aria2c AddURI error:%s", err)
		}
	} else {
		taskGID, err = aria2cRPCClient.AddTorrent(torrentBase64)
		if err != nil {
			return t.Errorf("call aria2c AddTorrent error:%s", err)
		}
		// save to file and change the sourceURL
		torrentFilename := fmt.Sprintf("%s/%s.torrent", downloadDir, torrentBase64[:16])
		fp, err := os.Create(torrentFilename)
		if err != nil {
			log.Warnf("os.Create file error:%s", err)
		}
		_, err = fp.Write(data)
		if err != nil {
			log.Warnf("fp.Write error:%s", err)
		}
		t.SourceURL = torrentFilename
	}
	log.Infof("create Magnet task: sourceURL:%s, taskGID:%s", t.SourceURL, taskGID)
	t.StartTime = time.Now()
	timeout := limitTimeout
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
MagnetLoop:
	complete := false
	ticker := time.NewTicker(time.Second * 5)
	for !complete {
		select {
		case <-ticker.C:
			result, err := aria2cRPCClient.TellStatus(taskGID)
			if err != nil {
				return t.Errorf("call aria2c TellStatus error:%s", err)
			}
			// update break condition
			complete = result.Completed()
			// udpate taskInfo
			t.TaskInfo.ContentLength = result.TotalLength
			t.Size = result.CompletedLength
			t.Duration = time.Now().Sub(t.StartTime)
			t.Speed = result.DownloadSpeed
			if !complete && result.CompletedLength > 0 && result.CompletedLength == result.TotalLength {
				// why aria2c wait so long time even it seems the download task is completed, may be as a seeder?
				log.Infof("force to set task status complete, status:%+v", result)
				time.Sleep(time.Second * 5)
				log.Infof("force to set task status complete, status:%+v", result)
				complete = true
			}
			// 磁力链接建立任务时无法指定文件名 获得真实文件名后需要重命名
			realFilename := strings.TrimPrefix(result.GetFilePath(), downloadDir+"/")
			pos := strings.Index(realFilename, "/")
			if pos != -1 && pos < len(realFilename)-1 {
				realFilename = realFilename[:pos]
			}
			t.TaskInfo.FileName = realFilename
			// 检查是否有继续下载磁力链接包含的其他文件
			followedBys := result.FollowedBy
			for _, followedTaskGID := range followedBys {
				taskGID = followedTaskGID
				complete = false
				log.Debugf("goto MagnetLoop: task status:%+v", result)
				goto MagnetLoop
			}
		case <-ctx.Done():
			if ctx.Err() == context.DeadlineExceeded {
				return t.Errorf("task timeout:%s", timeout)
			}
		}
	}
	err = aria2cRPCClient.RemoveDownloadResult(taskGID)
	if err != nil {
		log.Warnf("RemoveDownloadResult error:%s", err)
	}
	t.TaskInfo.IsCompleted = true
	t.Duration = time.Now().Sub(t.StartTime)
	t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
	return nil
}

func (t *MagnetTask) IsCompleted() bool {
	return t.TaskInfo.IsCompleted
}

func (t *MagnetTask) FileName() string {
	return t.TaskInfo.FileName
}

func (t *MagnetTask) ContentLength() int64 {
	return t.TaskInfo.ContentLength
}

func (t *MagnetTask) Errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	_, file, line, ok := runtime.Caller(1)
	if ok {
		log.Errorf("[%s:%d]%s", file, line, err.Error())
	}
	t.TaskInfo.IsError = true
	t.TaskInfo.IsCompleted = true
	t.TaskInfo.Error = err.Error()
	return err
}

//utils
var safeFilenameRegexp = regexp.MustCompile(`[\w\d\-.]+`)

func getSafeFilename(str string) string {
	filename := strings.Join(safeFilenameRegexp.FindAllString(str, -1), "")
	if lenOfFilename := len(filename); lenOfFilename > 50 {
		filename = filename[lenOfFilename-50 : lenOfFilename]
	}
	return fmt.Sprintf("%d-%s", time.Now().Unix(), filename)
}

func getHumanSizeString(byteSize int64) string {
	units := []string{"B", "KB", "MB", "GB", "TB", "EB"}
	index := 0
	byteSizeFloat := float64(byteSize)
	for ; byteSizeFloat > 1024; index += 1 {
		byteSizeFloat /= 1024
	}
	var unit string
	if index < len(units) {
		unit = units[index]
	} else {
		unit = "INF"
	}
	return fmt.Sprintf("%.2f %s", byteSizeFloat, unit)
}

func calculateDownloadSpeed(byteSize int64, duration time.Duration) (bytePerSecond int64) {
	if duration <= 0 {
		duration = 1
	}
	return byteSize * 1e9 / int64(duration)
}

func IsBase64String(str string) bool {
	_, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"fmt"
	"log"
	"os"
	"sync"
)

var wg sync.WaitGroup

const NUM_WORKERS = 10

//doDownload(host, query, destination beginYear, endYear, dataLimit)
func doDownload(host, where, destination, timeunit string, beginYear, endYear, dataLimit int) {
	uuidQuery := fmt.Sprintf("select distinct uuid where %s;", where)
	uuids := doDistinctQuery(host, uuidQuery)

	uuidQueue := make(chan string, NUM_WORKERS)

	wg.Add(len(uuids))
	for w := 0; w < NUM_WORKERS; w++ {
		go runWorker(uuidQueue, host, destination, timeunit, beginYear, endYear, dataLimit)
	}

	for _, uuid := range uuids {
		uuidQueue <- uuid
	}
	close(uuidQueue)
	wg.Wait()
}

func runWorker(uuids chan string, host, destination, timeunit string, beginYear, endYear, dataLimit int) {
	for uuid := range uuids {
		dlUUID(host, uuid, destination, timeunit, beginYear, endYear, dataLimit)
		wg.Done()
	}
}

func dlUUID(host, uuid, destination, timeunit string, beginYear, endYear, dataLimit int) {
	// open file
	filename := fmt.Sprintf("%s/%s.csv", destination, uuid)
	file, err := os.Create(filename)
	if err != nil {
		panic(err)
	}
	writer := csv.NewWriter(file)

	query := fmt.Sprintf("select data in ('1/1/%d', '1/1/%d') limit %d as %s where uuid = '%s';", endYear, beginYear, dataLimit, timeunit, uuid)
	msgs := doQuery(host, query)

	for _, msg := range msgs {
		numbers := toString(extractDataNumeric(msg))
		times := toString(extractTime(msg))
		log.Printf("%s -- %d points\n", msg.UUID, len(numbers))
		for i := 0; i < len(numbers); i++ {
			writer.Write([]string{times[i], numbers[i]})
		}
	}
	writer.Flush()
	file.Sync()
	file.Close()
}
<commit_msg>check errors<commit_after>package main

import (
	"encoding/csv"
	"fmt"
	"log"
	"os"
	"sync"
)

var wg sync.WaitGroup

const NUM_WORKERS = 10

//doDownload(host, query, destination beginYear, endYear, dataLimit)
func doDownload(host, where, destination, timeunit string, beginYear, endYear, dataLimit int) {
	uuidQuery := fmt.Sprintf("select distinct uuid where %s;", where)
	uuids := doDistinctQuery(host, uuidQuery)

	uuidQueue := make(chan string, NUM_WORKERS)

	wg.Add(len(uuids))
	for w := 0; w < NUM_WORKERS; w++ {
		go runWorker(uuidQueue, host, destination, timeunit, beginYear, endYear, dataLimit)
	}

	for _, uuid := range uuids {
		uuidQueue <- uuid
	}
	close(uuidQueue)
	wg.Wait()
}

func runWorker(uuids chan string, host, destination, timeunit string, beginYear, endYear, dataLimit int) {
	for uuid := range uuids {
		dlUUID(host, uuid, destination, timeunit, beginYear, endYear, dataLimit)
		wg.Done()
	}
}

func dlUUID(host, uuid, destination, timeunit string, beginYear, endYear, dataLimit int) {
	// open file
	filename := fmt.Sprintf("%s/%s.csv", destination, uuid)
	file, err := os.Create(filename)
	if err != nil {
		panic(err)
	}
	writer := csv.NewWriter(file)

	query := fmt.Sprintf("select data in ('1/1/%d', '1/1/%d') limit %d as %s where uuid = '%s';", endYear, beginYear, dataLimit, timeunit, uuid)
	msgs := doQuery(host, query)

	for _, msg := range msgs {
		numbers := toString(extractDataNumeric(msg))
		times := toString(extractTime(msg))
		log.Printf("%s -- %d points\n", msg.UUID, len(numbers))
		for i := 0; i < len(numbers); i++ {
			writer.Write([]string{times[i], numbers[i]})
		}
	}
	writer.Flush()
	if err := file.Sync(); err != nil {
		panic(err)
	}
	if err := file.Close(); err != nil {
		panic(err)
	}
}
<|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)
	bytesToSkip := -1

	switch {
	case t == "bool":
		bytesToSkip = 1
		break
	case t == "type" || t == "long" || t == "shor":
		bytesToSkip = 4
		break
	case t == "comp" || t == "dutc":
		bytesToSkip = 8
		break
	case t == "blob":
		blen, err := block.readUint32()
		if err != nil {
			break
		}
		bytesToSkip = int(blen)
		break
	case t == "ustr":
		blen, err := block.readUint32()
		if err != nil {
			break
		}
		bytesToSkip = int(2 * blen)
	default:
		break
	}

	if bytesToSkip <= 0 {
		return "", errors.New("Unknown file format")
	}
	block.skip(uint32(bytesToSkip))

	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 len(a.Offsets) <= int(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) {
	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>Should fix out-of-bounds slicing<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) {
	if len(a.Data) < pos+0x4+size {
		return nil, errors.New("Not enought Data")
	}
	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)
	bytesToSkip := -1

	switch {
	case t == "bool":
		bytesToSkip = 1
		break
	case t == "type" || t == "long" || t == "shor":
		bytesToSkip = 4
		break
	case t == "comp" || t == "dutc":
		bytesToSkip = 8
		break
	case t == "blob":
		blen, err := block.readUint32()
		if err != nil {
			break
		}
		bytesToSkip = int(blen)
		break
	case t == "ustr":
		blen, err := block.readUint32()
		if err != nil {
			break
		}
		bytesToSkip = int(2 * blen)
	default:
		break
	}

	if bytesToSkip <= 0 {
		return "", errors.New("Unknown file format")
	}
	block.skip(uint32(bytesToSkip))

	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, err = NewBlock(a, offset, size)
	if err != nil {
		return nil, err
	}

	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 len(a.Offsets) <= int(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) {
	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)
}
<|endoftext|>


<commit_before>package util

import (
	"encoding/json"
	"github.com/realglobe-Inc/go-lib-rg/erro"
	"time"
)

// "10s" みたいな形式で JSON にできる time.Duration のラッパー。

type Duration struct {
	time.Duration
}

func (dur *Duration) MarshalJSON() ([]byte, error) {
	return json.Marshal(dur.Duration.String())
}

func (dur *Duration) UnmarshalJSON(b []byte) error {
	var buff string
	err := json.Unmarshal(b, &buff)
	if err != nil {
		return erro.Wrap(err)
	}
	dur.Duration, err = time.ParseDuration(buff)
	if err != nil {
		return erro.Wrap(err)
	}
	return nil
}

func NewDuration(dur time.Duration) *Duration {
	return &Duration{dur}
}

func ParseDuration(s string) (*Duration, error) {
	dur, err := time.ParseDuration(s)
	if err != nil {
		return nil, erro.Wrap(err)
	}
	return &Duration{dur}, nil
}
<commit_msg>error 変数のスコープができるだけ小さくなるように変更<commit_after>package util

import (
	"encoding/json"
	"github.com/realglobe-Inc/go-lib-rg/erro"
	"time"
)

// "10s" みたいな形式で JSON にできる time.Duration のラッパー。

type Duration struct {
	time.Duration
}

func (dur *Duration) MarshalJSON() ([]byte, error) {
	return json.Marshal(dur.Duration.String())
}

func (dur *Duration) UnmarshalJSON(b []byte) error {
	var buff string
	if err := json.Unmarshal(b, &buff); err != nil {
		return erro.Wrap(err)
	}
	var err error
	dur.Duration, err = time.ParseDuration(buff)
	if err != nil {
		return erro.Wrap(err)
	}
	return nil
}

func NewDuration(dur time.Duration) *Duration {
	return &Duration{dur}
}

func ParseDuration(s string) (*Duration, error) {
	dur, err := time.ParseDuration(s)
	if err != nil {
		return nil, erro.Wrap(err)
	}
	return &Duration{dur}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

const (
	corednsImage = "k8s.gcr.io/coredns:1.6.2"
	testerImage  = "c.txtdirect.org/tester:dirty"
)

var txtdirectImage = fmt.Sprintf("c.txtdirect.org/txtdirect:%s", os.Getenv("VERSION"))

var resultRegex = regexp.MustCompile("Total:+\\s(\\d+),\\sPassed:+\\s(\\d+),\\sFailed:+\\s(\\d+)")

type dockerManager struct {
	dir       string
	cdir      string
	gomodpath string

	testerContainerLogs map[string][]byte

	stats struct {
		total  int
		passed int
		failed int
	}
}

func main() {
	d := dockerManager{
		testerContainerLogs: make(map[string][]byte),
	}

	if err := d.CheckClient(); err != nil {
		log.Fatalf("[txtdirect_e2e]: Docker daemon didn't respond to client: %s", err)
	}

	var directories []string
	if err := listDirectories(&directories); err != nil {
		log.Printf("[txtdirect_e2e]: Couldn't list the test directories: %s", err.Error())
	}

	if err := d.CreateNetwork(); err != nil {
		log.Fatalf("[txtdirect_e2e]: Couldn't create the network adaptor: %s", err.Error())
	}

	// Run the tests for each test-case
	for _, directory := range directories {
		// Start the CoreDNS and TXTDirect containers for test-case
		d.dir = directory
		if err := d.StartContainers(); err != nil {
			log.Fatalf("[txtdirect_e2e]: Couldn't start containers: %s", err.Error())
		}

		if err := d.RunTesterContainer(); err != nil {
			log.Fatalf("[txtdirect_e2e]: Couldn't run the tests in container: %s", err.Error())
		}

		if err := d.StopContainers(); err != nil {
			log.Fatalf("[txtdirect_e2e]: Couldn't stop containers: %s", err.Error())
		}
	}

	if err := d.ExamineLogs(); err != nil {
		log.Fatalf("[txtdirect_e2e]: Couldn't examine the logs and count the stats: %s", err.Error())
	}

	if err := d.RemoveNetwork(); err != nil {
		log.Fatalf("[txtdirect_e2e]: Couldn't remove the network adaptor: %s", err.Error())
	}

	if d.stats.failed != 0 {
		log.Fatalf("[txtdirect_e2e]: Tests failed because there were %d failed tests.", d.stats.failed)
	}
}

func listDirectories(directories *[]string) error {
	return filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return fmt.Errorf("Couldn't find the directory %q: %s", path, err.Error())
		}
		if info.IsDir() && path != "." {
			*directories = append(*directories, path)
			return filepath.SkipDir
		}
		return nil
	})
}

// CheckClient calls the "docker info" command to check if Docker daemon and CLI exist
func (d *dockerManager) CheckClient() error {
	_, err := exec.Command("docker", "info").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Cannot connect to the Docker daemon at unix:///var/run/docker.sock. Is the docker daemon running?")
	}

	return nil
}

// StartContainers starts a CoreDNS and a TXTDirect container for each test-case
func (d *dockerManager) StartContainers() error {
	// Get current working directory to create mountpoint to containers for test-case data
	var err error
	d.cdir, err = os.Getwd()
	if err != nil {
		return fmt.Errorf("Couldn't get the current working directory: %s", err.Error())
	}

	// Fetch GOPATH for mounting /go/pkg/mod volumes
	if os.Getenv("GOPATH") == "" {
		return fmt.Errorf("$GOPATH is empty")
	}
	d.gomodpath = fmt.Sprintf("%s/pkg/mod", os.Getenv("GOPATH"))

	// Create the CoreDNS container
	_, err = exec.Command("docker",
		"container", "run",
		"-d",
		// Exposed Ports
		"-p", "53:53",
		// Mounted volumes
		"-v", fmt.Sprintf("%s/%s:/e2e", d.cdir, d.dir),
		"-v", fmt.Sprintf("%s:/go/pkg/mod", d.gomodpath),
		// Network
		"--network", "coretxtd",
		// Container Name
		"--name", "e2e_coredns_container",
		corednsImage,
		// CMD
		"-conf",
		"/e2e/Corefile",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't create the CoreDNS container: %s", err.Error())
	}

	// Create the TXTDirect container
	_, err = exec.Command("docker",
		"container", "run",
		"-d",
		// Exposed Ports
		"-p", "80:80",
		// Mounted volumes
		"-v", fmt.Sprintf("%s/%s:/e2e", d.cdir, d.dir),
		"-v", fmt.Sprintf("%s:/go/pkg/mod", d.gomodpath),
		// Network
		"--network", "coretxtd",
		// Container Name
		"--name", "e2e_txtdirect_container",
		txtdirectImage,
		// CMD
		"-conf",
		"/e2e/txtdirect.config",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't create the TXTDirect container: %s", err.Error())
	}

	// Connect the Docker registry to the e2e network
	if strings.Contains(d.dir, "dockerv2") {
		_, err := exec.Command("docker", "network", "connect", "coretxtd", "registry").CombinedOutput()
		if err != nil {
			return fmt.Errorf("Couldn't connect Docker registry the network adaptor: %s", err.Error())
		}

		// Tag TXTDirect's image to use in custom registry
		_, err = exec.Command("docker", "tag", txtdirectImage, "172.20.10.3:5000/txtdirect").CombinedOutput()
		if err != nil {
			return fmt.Errorf("Couldn't tag image to use in custom Docker registry: %s", err.Error())
		}

		_, err = exec.Command("docker", "save", "172.20.10.3:5000/txtdirect:latest", "-o", "txtdirect.tar").CombinedOutput()
		if err != nil {
			return fmt.Errorf("Couldn't export the image into a tarball: %s", err.Error())
		}

		// Push the TXTDirect image to the custom registry
		_, err = exec.Command("crane", "push", "txtdirect.tar", "172.20.10.3:5000/txtdirect").CombinedOutput()
		if err != nil {
			return fmt.Errorf("Couldn't push the image to the custom Docker registry: %s", err.Error())
		}

		_, err = exec.Command("rm", "./txtdirect.tar").CombinedOutput()
		if err != nil {
			return fmt.Errorf("Couldn't remove the image tarball: %s", err.Error())
		}
	}

	return nil
}

func (d *dockerManager) StopContainers() error {
	if strings.Contains(d.dir, "dockerv2") {
		_, err := exec.Command("docker",
			"logs",
			"e2e_txtdirect_container",
		).CombinedOutput()
		if err != nil {
			return err
		}
	}
	_, err := exec.Command("docker",
		"container", "rm", "-f",
		"e2e_coredns_container",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the CoreDNS container: %s", err.Error())
	}

	_, err = exec.Command("docker",
		"container", "rm", "-f",
		"e2e_txtdirect_container",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the TXTDirect container: %s", err.Error())
	}

	_, err = exec.Command("docker",
		"container", "rm", "-f",
		fmt.Sprintf("e2e_%s_container", d.dir),
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the tester container: %s", err.Error())
	}

	if strings.Contains(d.dir, "dockerv2") {
		_, err = exec.Command("docker",
			"network", "disconnect",
			"coretxtd",
			"registry",
		).CombinedOutput()
		if err != nil {
			return fmt.Errorf("Couldn't disconnect the Docker registry container: %s", err.Error())
		}
	}

	return nil
}

func (d *dockerManager) RunTesterContainer() error {
	_, err := user.Current()
	if err != nil {
		return fmt.Errorf("Couldn't get the current user: %s", err.Error())
	}

	// Create the tester container
	d.testerContainerLogs[d.dir], err = exec.Command("docker",
		"container", "run",
		// Mounted volumes
		"-v", fmt.Sprintf("%s/%s:/e2e", d.cdir, d.dir),
		"-v", fmt.Sprintf("%s:/go/pkg/mod", d.gomodpath),
		// DNS
		"--dns", "172.20.10.1",
		// Network
		"--network", "coretxtd",
		// Container Name
		"--name", fmt.Sprintf("e2e_%s_container", d.dir),
		testerImage,
		// CMD
		"go", "run", "main.go",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't create the tester container: %s", err.Error())
	}

	return nil
}

// ExamineLogs reads each tester container's logs and counts the failed and passed tests
func (d *dockerManager) ExamineLogs() error {
	for testCase, logs := range d.testerContainerLogs {
		log.Println(string(logs))
		// Stats slice order: [0: Full Log, 1: Total, 2: Passed, 3: Failed]
		stats := resultRegex.FindAllStringSubmatch(string(logs), -1)[0]
		if err := d.CountStats(stats); err != nil {
			return fmt.Errorf("Couldn't count the stats for \"%s\" test-case: %s", testCase, err.Error())
		}
	}
	log.Printf("Total Tests: %d, Total Passed Tests: %d, Total Failed Tests: %d", d.stats.total, d.stats.passed, d.stats.failed)
	return nil
}

func (d *dockerManager) CountStats(stats []string) error {
	total, err := strconv.Atoi(stats[1])
	if err != nil {
		return fmt.Errorf("Couldn't conver count of tests to int: %s", err.Error())
	}
	d.stats.total += total

	passed, err := strconv.Atoi(stats[2])
	if err != nil {
		return fmt.Errorf("Couldn't convert count of passed tests to int: %s", err.Error())
	}
	d.stats.passed += passed

	failed, err := strconv.Atoi(stats[3])
	if err != nil {
		return fmt.Errorf("Couldn't convert count of failed tests to int: %s", err.Error())
	}
	d.stats.failed += failed

	return nil
}

func (d *dockerManager) RemoveNetwork() error {
	_, err := exec.Command("docker",
		"network", "rm",
		"coretxtd",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the network adaptor: %s", err.Error())
	}
	return nil
}

func (d *dockerManager) CreateNetwork() error {
	// Create a Docker network for containers
	_, err := exec.Command("docker", "network", "create", "--ip-range=\"172.20.10.0/24\"", "--subnet=\"172.20.0.0/16\"", "coretxtd").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't create the network adaptor: %s", err.Error())
	}
	return nil
}
<commit_msg>(e2e): Move test image push to a separate method<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

const (
	corednsImage = "k8s.gcr.io/coredns:1.6.2"
	testerImage  = "c.txtdirect.org/tester:dirty"
)

var txtdirectImage = fmt.Sprintf("c.txtdirect.org/txtdirect:%s", os.Getenv("VERSION"))

var resultRegex = regexp.MustCompile("Total:+\\s(\\d+),\\sPassed:+\\s(\\d+),\\sFailed:+\\s(\\d+)")

type dockerManager struct {
	dir       string
	cdir      string
	gomodpath string

	testerContainerLogs map[string][]byte

	stats struct {
		total  int
		passed int
		failed int
	}
}

func main() {
	d := dockerManager{
		testerContainerLogs: make(map[string][]byte),
	}

	if err := d.CheckClient(); err != nil {
		log.Fatalf("[txtdirect_e2e]: Docker daemon didn't respond to client: %s", err)
	}

	var directories []string
	if err := listDirectories(&directories); err != nil {
		log.Printf("[txtdirect_e2e]: Couldn't list the test directories: %s", err.Error())
	}

	if err := d.CreateNetwork(); err != nil {
		log.Fatalf("[txtdirect_e2e]: Couldn't create the network adaptor: %s", err.Error())
	}

	// Run the tests for each test-case
	for _, directory := range directories {
		// Start the CoreDNS and TXTDirect containers for test-case
		d.dir = directory

		if err := d.StartContainers(); err != nil {
			log.Fatalf("[txtdirect_e2e]: Couldn't start containers: %s", err.Error())
		}

		if err := d.RunTesterContainer(); err != nil {
			log.Fatalf("[txtdirect_e2e]: Couldn't run the tests in container: %s", err.Error())
		}

		if err := d.StopContainers(); err != nil {
			log.Fatalf("[txtdirect_e2e]: Couldn't stop containers: %s", err.Error())
		}
	}

	if err := d.ExamineLogs(); err != nil {
		log.Fatalf("[txtdirect_e2e]: Couldn't examine the logs and count the stats: %s", err.Error())
	}

	if err := d.RemoveNetwork(); err != nil {
		log.Fatalf("[txtdirect_e2e]: Couldn't remove the network adaptor: %s", err.Error())
	}

	if d.stats.failed != 0 {
		log.Fatalf("[txtdirect_e2e]: Tests failed because there were %d failed tests.", d.stats.failed)
	}
}

func listDirectories(directories *[]string) error {
	return filepath.Walk(".", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return fmt.Errorf("Couldn't find the directory %q: %s", path, err.Error())
		}
		if info.IsDir() && path != "." {
			*directories = append(*directories, path)
			return filepath.SkipDir
		}
		return nil
	})
}

// CheckClient calls the "docker info" command to check if Docker daemon and CLI exist
func (d *dockerManager) CheckClient() error {
	_, err := exec.Command("docker", "info").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Cannot connect to the Docker daemon at unix:///var/run/docker.sock. Is the docker daemon running?")
	}

	return nil
}

// StartContainers starts a CoreDNS and a TXTDirect container for each test-case
func (d *dockerManager) StartContainers() error {
	// Get current working directory to create mountpoint to containers for test-case data
	var err error
	d.cdir, err = os.Getwd()
	if err != nil {
		return fmt.Errorf("Couldn't get the current working directory: %s", err.Error())
	}

	// Fetch GOPATH for mounting /go/pkg/mod volumes
	if os.Getenv("GOPATH") == "" {
		return fmt.Errorf("$GOPATH is empty")
	}
	d.gomodpath = fmt.Sprintf("%s/pkg/mod", os.Getenv("GOPATH"))

	// Create the CoreDNS container
	_, err = exec.Command("docker",
		"container", "run",
		"-d",
		// Exposed Ports
		"-p", "53:53",
		// Mounted volumes
		"-v", fmt.Sprintf("%s/%s:/e2e", d.cdir, d.dir),
		"-v", fmt.Sprintf("%s:/go/pkg/mod", d.gomodpath),
		// Network
		"--network", "coretxtd",
		// Container Name
		"--name", "e2e_coredns_container",
		corednsImage,
		// CMD
		"-conf",
		"/e2e/Corefile",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't create the CoreDNS container: %s", err.Error())
	}

	// Create the TXTDirect container
	_, err = exec.Command("docker",
		"container", "run",
		"-d",
		// Exposed Ports
		"-p", "80:80",
		// Mounted volumes
		"-v", fmt.Sprintf("%s/%s:/e2e", d.cdir, d.dir),
		"-v", fmt.Sprintf("%s:/go/pkg/mod", d.gomodpath),
		// Network
		"--network", "coretxtd",
		// Container Name
		"--name", "e2e_txtdirect_container",
		txtdirectImage,
		// CMD
		"-conf",
		"/e2e/txtdirect.config",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't create the TXTDirect container: %s", err.Error())
	}

	// Push the test image to the testing Docker registry
	if strings.Contains(d.dir, "dockerv2") {
		if err := d.PushTestImage(); err != nil {
			return err
		}
	}

	return nil
}

func (d *dockerManager) StopContainers() error {
	if strings.Contains(d.dir, "dockerv2") {
		_, err := exec.Command("docker",
			"logs",
			"e2e_txtdirect_container",
		).CombinedOutput()
		if err != nil {
			return err
		}
	}
	_, err := exec.Command("docker",
		"container", "rm", "-f",
		"e2e_coredns_container",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the CoreDNS container: %s", err.Error())
	}

	_, err = exec.Command("docker",
		"container", "rm", "-f",
		"e2e_txtdirect_container",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the TXTDirect container: %s", err.Error())
	}

	_, err = exec.Command("docker",
		"container", "rm", "-f",
		fmt.Sprintf("e2e_%s_container", d.dir),
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the tester container: %s", err.Error())
	}

	if strings.Contains(d.dir, "dockerv2") {
		_, err = exec.Command("docker",
			"network", "disconnect",
			"coretxtd",
			"registry",
		).CombinedOutput()
		if err != nil {
			return fmt.Errorf("Couldn't disconnect the Docker registry container: %s", err.Error())
		}
	}

	return nil
}

func (d *dockerManager) RunTesterContainer() error {
	_, err := user.Current()
	if err != nil {
		return fmt.Errorf("Couldn't get the current user: %s", err.Error())
	}

	// Create the tester container
	d.testerContainerLogs[d.dir], err = exec.Command("docker",
		"container", "run",
		// Mounted volumes
		"-v", fmt.Sprintf("%s/%s:/e2e", d.cdir, d.dir),
		"-v", fmt.Sprintf("%s:/go/pkg/mod", d.gomodpath),
		// DNS
		"--dns", "172.20.10.1",
		// Network
		"--network", "coretxtd",
		// Container Name
		"--name", fmt.Sprintf("e2e_%s_container", d.dir),
		testerImage,
		// CMD
		"go", "run", "main.go",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't create the tester container: %s", err.Error())
	}

	return nil
}

// ExamineLogs reads each tester container's logs and counts the failed and passed tests
func (d *dockerManager) ExamineLogs() error {
	for testCase, logs := range d.testerContainerLogs {
		log.Println(string(logs))
		// Stats slice order: [0: Full Log, 1: Total, 2: Passed, 3: Failed]
		stats := resultRegex.FindAllStringSubmatch(string(logs), -1)[0]
		if err := d.CountStats(stats); err != nil {
			return fmt.Errorf("Couldn't count the stats for \"%s\" test-case: %s", testCase, err.Error())
		}
	}
	log.Printf("Total Tests: %d, Total Passed Tests: %d, Total Failed Tests: %d", d.stats.total, d.stats.passed, d.stats.failed)
	return nil
}

func (d *dockerManager) CountStats(stats []string) error {
	total, err := strconv.Atoi(stats[1])
	if err != nil {
		return fmt.Errorf("Couldn't conver count of tests to int: %s", err.Error())
	}
	d.stats.total += total

	passed, err := strconv.Atoi(stats[2])
	if err != nil {
		return fmt.Errorf("Couldn't convert count of passed tests to int: %s", err.Error())
	}
	d.stats.passed += passed

	failed, err := strconv.Atoi(stats[3])
	if err != nil {
		return fmt.Errorf("Couldn't convert count of failed tests to int: %s", err.Error())
	}
	d.stats.failed += failed

	return nil
}

func (d *dockerManager) RemoveNetwork() error {
	_, err := exec.Command("docker",
		"network", "rm",
		"coretxtd",
	).CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the network adaptor: %s", err.Error())
	}
	return nil
}

func (d *dockerManager) CreateNetwork() error {
	// Create a Docker network for containers
	_, err := exec.Command("docker", "network", "create", "--ip-range=\"172.20.10.0/24\"", "--subnet=\"172.20.0.0/16\"", "coretxtd").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't create the network adaptor: %s", err.Error())
	}
	return nil
}

// PushTestImage connects the registry to network adaptor and pushes
// a testing image to the registry
func (d *dockerManager) PushTestImage() error {
	// Connect registry to the network adaptor
	_, err := exec.Command("docker", "network", "connect", "coretxtd", "registry").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't connect Docker registry the network adaptor: %s", err.Error())
	}

	// Tag TXTDirect's image to use in custom registry
	_, err = exec.Command("docker", "tag", txtdirectImage, "172.20.10.3:5000/txtdirect").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't tag image to use in custom Docker registry: %s", err.Error())
	}

	// Export the image to a tarball
	_, err = exec.Command("docker", "save", "172.20.10.3:5000/txtdirect:latest", "-o", "txtdirect.tar").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't export the image into a tarball: %s", err.Error())
	}

	// Push the TXTDirect image to the custom registry
	_, err = exec.Command("crane", "push", "txtdirect.tar", "172.20.10.3:5000/txtdirect").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't push the image to the custom Docker registry: %s", err.Error())
	}

	// Remove the tarball after the push
	_, err = exec.Command("rm", "./txtdirect.tar").CombinedOutput()
	if err != nil {
		return fmt.Errorf("Couldn't remove the image tarball: %s", err.Error())
	}

	return nil
}
<|endoftext|>


<commit_before>package gorpc

import (
	"bufio"
	"compress/flate"
	"encoding/gob"
	"io"
)

// RegisterType registers the given type to send via rpc.
//
// The client must register all the response types the server may send.
// The server must register all the request types the client may send.
//
// There is no need in registering base Go types such as int, string, bool,
// float64, etc. or arrays, slices and maps containing base Go types.
func RegisterType(x interface{}) {
	gob.Register(x)
}

type wireRequest struct {
	ID           uint64
	Request      interface{}
	SkipResponse bool
}

type wireResponse struct {
	ID       uint64
	Response interface{}
	Error    string
}

type messageEncoder struct {
	e  *gob.Encoder
	bw *bufio.Writer
	zw *flate.Writer
	ww *bufio.Writer
}

func (e *messageEncoder) Close() error {
	if e.zw != nil {
		return e.zw.Close()
	}
	return nil
}

func (e *messageEncoder) Flush() error {
	if e.zw != nil {
		if err := e.ww.Flush(); err != nil {
			return err
		}
		if err := e.zw.Flush(); err != nil {
			return err
		}
	}
	if err := e.bw.Flush(); err != nil {
		return err
	}
	return nil
}

func (e *messageEncoder) Encode(msg interface{}) error {
	return e.e.Encode(msg)
}

func newMessageEncoder(w io.Writer, bufferSize int, enableCompression bool, s *ConnStats) *messageEncoder {
	w = newWriterCounter(w, s)
	bw := bufio.NewWriterSize(w, bufferSize)

	ww := bw
	var zw *flate.Writer
	if enableCompression {
		zw, _ = flate.NewWriter(bw, flate.BestSpeed)
		ww = bufio.NewWriterSize(zw, bufferSize)
	}

	return &messageEncoder{
		e:  gob.NewEncoder(ww),
		bw: bw,
		zw: zw,
		ww: ww,
	}
}

type messageDecoder struct {
	d  *gob.Decoder
	zr io.ReadCloser
}

func (d *messageDecoder) Close() error {
	if d.zr != nil {
		return d.zr.Close()
	}
	return nil
}

func (d *messageDecoder) Decode(msg interface{}) error {
	return d.d.Decode(msg)
}

func newMessageDecoder(r io.Reader, bufferSize int, enableCompression bool, s *ConnStats) *messageDecoder {
	r = newReaderCounter(r, s)
	br := bufio.NewReaderSize(r, bufferSize)

	rr := br
	var zr io.ReadCloser
	if enableCompression {
		zr = flate.NewReader(br)
		rr = bufio.NewReaderSize(zr, bufferSize)
	}

	return &messageDecoder{
		d:  gob.NewDecoder(rr),
		zr: zr,
	}
}
<commit_msg>Subsituted compress/flate with github.com/pierrec/lz4<commit_after>package gorpc

import (
	"bufio"
	"encoding/gob"
	"github.com/pierrec/lz4"
	"io"
)

// RegisterType registers the given type to send via rpc.
//
// The client must register all the response types the server may send.
// The server must register all the request types the client may send.
//
// There is no need in registering base Go types such as int, string, bool,
// float64, etc. or arrays, slices and maps containing base Go types.
func RegisterType(x interface{}) {
	gob.Register(x)
}

type wireRequest struct {
	ID           uint64
	Request      interface{}
	SkipResponse bool
}

type wireResponse struct {
	ID       uint64
	Response interface{}
	Error    string
}

type messageEncoder struct {
	e  *gob.Encoder
	bw *bufio.Writer
	zw *lz4.Writer
	ww *bufio.Writer
}

func (e *messageEncoder) Close() error {
	if e.zw != nil {
		return e.zw.Close()
	}
	return nil
}

func (e *messageEncoder) Flush() error {
	if e.zw != nil {
		if err := e.ww.Flush(); err != nil {
			return err
		}
		if err := e.zw.Flush(); err != nil {
			return err
		}
	}
	if err := e.bw.Flush(); err != nil {
		return err
	}
	return nil
}

func (e *messageEncoder) Encode(msg interface{}) error {
	return e.e.Encode(msg)
}

func newMessageEncoder(w io.Writer, bufferSize int, enableCompression bool, s *ConnStats) *messageEncoder {
	w = newWriterCounter(w, s)
	bw := bufio.NewWriterSize(w, bufferSize)

	ww := bw
	var zw *lz4.Writer
	if enableCompression {
		zw = lz4.NewWriter(bw)
		ww = bufio.NewWriterSize(zw, bufferSize)
	}

	return &messageEncoder{
		e:  gob.NewEncoder(ww),
		bw: bw,
		zw: zw,
		ww: ww,
	}
}

type messageDecoder struct {
	d *gob.Decoder
}

func (d *messageDecoder) Close() error {
	return nil
}

func (d *messageDecoder) Decode(msg interface{}) error {
	return d.d.Decode(msg)
}

func newMessageDecoder(r io.Reader, bufferSize int, enableCompression bool, s *ConnStats) *messageDecoder {
	r = newReaderCounter(r, s)
	br := bufio.NewReaderSize(r, bufferSize)

	rr := br
	if enableCompression {
		zr := lz4.NewReader(br)
		rr = bufio.NewReaderSize(zr, bufferSize)
	}

	return &messageDecoder{
		d: gob.NewDecoder(rr),
	}
}
<|endoftext|>