<commit_before>package web

import (
	"errors"
	"net/http"
	"net/url"
	"time"

	log "github.com/Ptt-Alertor/logrus"

	"github.com/Ptt-Alertor/ptt-alertor/models/article"

	"regexp"

	"strings"

	"strconv"

	"github.com/Ptt-Alertor/ptt-alertor/models/pushsum"
	"golang.org/x/net/html"
)

const pttHostURL = "https://www.ptt.cc"

// CurrentPage find Board Last Page Number
func CurrentPage(board string) (int, error) {
	url := makeBoardURL(board, -1)
	htmlNodes, err := fetchHTML(url)
	if err != nil {
		return 0, err
	}
	paging := findNodes(htmlNodes, findPagingBlock)
	for _, page := range paging {
		anchors := findNodes(page, findAnchor)
		for _, a := range anchors {
			if strings.Contains(a.FirstChild.Data, "上頁") {
				link := getAnchorLink(a)
				re := regexp.MustCompile("\\d+")
				page, err := strconv.Atoi(re.FindString(link))
				if err != nil {
					return 0, err
				}
				return page + 1, nil
			}
		}
	}
	return 0, errors.New("Parse Currenect Page Error")
}

// FetchArticles makes board's index articles to a article slice
func FetchArticles(board string, page int) (articles article.Articles, err error) {
	reqURL := makeBoardURL(board, page)
	htmlNodes, err := fetchHTML(reqURL)
	if err != nil {
		return nil, err
	}
	articleBlocks := findNodes(htmlNodes, findArticleBlocks)
	for _, articleBlock := range articleBlocks {
		article := article.Article{}
		for _, pushCountDiv := range findNodes(articleBlock, findPushCountDiv) {
			if child := pushCountDiv.FirstChild; child != nil {
				if child := child.FirstChild; child != nil {
					if err == nil {
						article.PushSum = pushsum.ConvertPushCount(child.Data)
					}
				}
			}
		}
		for _, titleDiv := range findNodes(articleBlock, findTitleDiv) {

			anchors := findNodes(titleDiv, findAnchor)

			if len(anchors) == 0 {
				article.Title = titleDiv.FirstChild.Data
				article.Link = ""
				continue
			}

			for _, anchor := range findNodes(titleDiv, findAnchor) {
				article.Title = anchor.FirstChild.Data
				link := pttHostURL + getAnchorLink(anchor)
				article.Link = link
				article.ID = article.ParseID(link)
			}
		}
		for _, metaDiv := range findNodes(articleBlock, findMetaDiv) {
			for _, date := range findNodes(metaDiv, findDateDiv) {
				article.Date = strings.TrimSpace(date.FirstChild.Data)
			}
			for _, author := range findNodes(metaDiv, findAuthorDiv) {
				article.Author = author.FirstChild.Data
			}
		}
		articles = append(articles, article)
		if isLastArticleBlock(articleBlock) {
			break
		}
	}
	return articles, nil
}

func isLastArticleBlock(articleBlock *html.Node) bool {
	for next := articleBlock.NextSibling; ; next = next.NextSibling {
		if next == nil {
			break
		}
		if next.Type == html.ElementNode {
			for _, attr := range next.Attr {
				if attr.Val == "r-list-sep" {
					return true
				}
				return false
			}
		}
	}
	return false
}

// FetchArticle build article object from html
func FetchArticle(board, articleCode string) (article.Article, error) {
	reqURL := makeArticleURL(board, articleCode)
	htmlNodes, err := fetchHTML(reqURL)
	if err != nil {
		return article.Article{}, err
	}
	atcl := article.Article{
		Link:  reqURL,
		Code:  articleCode,
		Board: board,
	}
	nodes := findNodes(htmlNodes, findOgTitleMeta)
	if len(nodes) > 0 {
		atcl.Title = getMetaContent(nodes[0])
	} else {
		atcl.Title = "[內文標題已被刪除]"
	}
	atcl.ID = atcl.ParseID(reqURL)
	pushBlocks := findNodes(htmlNodes, findPushBlocks)
	pushes := []article.Comment{}

PushBlocks:
	for _, pushBlock := range pushBlocks {
		push := article.Comment{}
		for _, pushIPDateTime := range findNodes(pushBlock, findPushIPDateTime) {
			ipdatetime := strings.TrimSpace(pushIPDateTime.FirstChild.Data)
			if ipdatetime == "" {
				continue PushBlocks
			}
			dateTime, err := parseDateTime(ipdatetime)
			if err != nil {
				log.WithFields(log.Fields{
					"ipdatetime": ipdatetime,
					"board":      board,
					"code":       articleCode,
				}).WithError(err).Warn("Parse DateTime Error")
				continue PushBlocks
			}
			push.DateTime = dateTime
		}
		for _, pushTag := range findNodes(pushBlock, findPushTag) {
			push.Tag = pushTag.FirstChild.Data
		}
		for _, pushUserID := range findNodes(pushBlock, findPushUserID) {
			push.UserID = pushUserID.FirstChild.Data
		}
		for _, pushContent := range findNodes(pushBlock, findPushContent) {
			content := pushContent.FirstChild.Data
			for n := pushContent.FirstChild.NextSibling; n != nil; n = n.NextSibling {
				if findEmailProtected(n) != nil {
					break
				}
				if n.FirstChild != nil {
					content += n.FirstChild.Data
				}
				if n.NextSibling != nil {
					content += n.NextSibling.Data
				}
			}
			push.Content = content
		}
		pushes = append(pushes, push)
	}
	if len(pushes) > 0 {
		atcl.LastPushDateTime = pushes[len(pushes)-1].DateTime
	}
	atcl.Comments = pushes
	return atcl, nil
}

func parseDateTime(ipdatetime string) (time.Time, error) {
	re, _ := regexp.Compile("(\\d+\\.\\d+\\.\\d+\\.\\d+)?\\s*(.*)")
	subMatches := re.FindStringSubmatch(ipdatetime)
	dateTime := strings.TrimSpace(subMatches[len(subMatches)-1])
	loc := time.FixedZone("CST", 8*60*60)
	t, err := time.ParseInLocation("01/02 15:04", dateTime, loc)
	if err != nil {
		return time.Time{}, err
	}
	t = t.AddDate(getYear(t), 0, 0)
	return t, nil
}

func getYear(commentTime time.Time) int {
	t := time.Now()
	commentTime = commentTime.AddDate(t.Year(), 0, 0)
	if commentTime.After(t) {
		return t.Year() - 1
	}
	return t.Year()
}

// CheckBoardExist use for checking board exist or not
func CheckBoardExist(board string) bool {
	return checkURLExist(makeBoardURL(board, -1))
}

// CheckArticleExist user for checking article exist or not
func CheckArticleExist(board, articleCode string) bool {
	return checkURLExist(makeArticleURL(board, articleCode))
}

func checkURLExist(url string) bool {
	resp, err := http.Get(url)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusOK {
		return true
	}
	return false
}

func makeBoardURL(board string, page int) string {
	var pageStr string
	if page < 0 {
		pageStr = ""
	} else {
		pageStr = strconv.Itoa(page)
	}
	return pttHostURL + "/bbs/" + board + "/index" + pageStr + ".html"
}

func makeArticleURL(board, articleCode string) string {
	return pttHostURL + "/bbs/" + board + "/" + articleCode + ".html"
}

// URLNotFoundError is an error type present 404 Not Found
type URLNotFoundError struct {
	URL string
}

func (u URLNotFoundError) Error() string {
	return "Fetched URL Not Found"
}

var errRedirect = errors.New("Redirect")

var client = &http.Client{
	CheckRedirect: func(req *http.Request, via []*http.Request) error {
		return errRedirect
	},
	Timeout: 30 * time.Second,
}

func fetchHTML(reqURL string) (doc *html.Node, err error) {
	resp, err := client.Get(reqURL)
	if err != nil && resp == nil {
		log.WithField("url", reqURL).WithError(err).Error("Fetch URL Failed")
		return nil, err
	}
	defer resp.Body.Close()

	if err == nil && resp.StatusCode == http.StatusNotFound {
		err = URLNotFoundError{reqURL}
		return nil, err
	}

	if uerr, ok := err.(*url.Error); ok && uerr.Err == errRedirect {
		req := passR18(reqURL)
		resp, err = client.Do(req)
		if err != nil {
			return nil, err
		}
	}

	doc, err = html.Parse(resp.Body)
	if err != nil {
		log.WithError(err).Error("Crawler Fetch HTML Failed")
		return nil, err
	}
	return doc, nil
}

func passR18(reqURL string) (req *http.Request) {

	req, _ = http.NewRequest("GET", reqURL, nil)

	over18Cookie := http.Cookie{
		Name:       "over18",
		Value:      "1",
		Domain:     "www.ptt.cc",
		Path:       "/",
		RawExpires: "Session",
		MaxAge:     0,
		HttpOnly:   false,
	}

	req.AddCookie(&over18Cookie)

	return req
}
<commit_msg>:bug: fix the board only has one page on CurrentPage<commit_after>package web

import (
	"errors"
	"net/http"
	"net/url"
	"time"

	log "github.com/Ptt-Alertor/logrus"

	"github.com/Ptt-Alertor/ptt-alertor/models/article"

	"regexp"

	"strings"

	"strconv"

	"github.com/Ptt-Alertor/ptt-alertor/models/pushsum"
	"golang.org/x/net/html"
)

const pttHostURL = "https://www.ptt.cc"

// CurrentPage find Board Last Page Number
func CurrentPage(board string) (int, error) {
	url := makeBoardURL(board, -1)
	htmlNodes, err := fetchHTML(url)
	if err != nil {
		return 0, err
	}
	paging := findNodes(htmlNodes, findPagingBlock)
	for _, page := range paging {
		anchors := findNodes(page, findAnchor)
		for _, a := range anchors {
			if strings.Contains(a.FirstChild.Data, "上頁") {
				link := getAnchorLink(a)
				if link == "" {
					return 1, nil
				}
				re := regexp.MustCompile("\\d+")
				page, err := strconv.Atoi(re.FindString(link))
				if err != nil {
					return 0, err
				}
				return page + 1, nil
			}
		}
	}
	return 0, errors.New("Parse Current Page Error")
}

// FetchArticles makes board's index articles to a article slice
func FetchArticles(board string, page int) (articles article.Articles, err error) {
	reqURL := makeBoardURL(board, page)
	htmlNodes, err := fetchHTML(reqURL)
	if err != nil {
		return nil, err
	}
	articleBlocks := findNodes(htmlNodes, findArticleBlocks)
	for _, articleBlock := range articleBlocks {
		article := article.Article{}
		for _, pushCountDiv := range findNodes(articleBlock, findPushCountDiv) {
			if child := pushCountDiv.FirstChild; child != nil {
				if child := child.FirstChild; child != nil {
					if err == nil {
						article.PushSum = pushsum.ConvertPushCount(child.Data)
					}
				}
			}
		}
		for _, titleDiv := range findNodes(articleBlock, findTitleDiv) {

			anchors := findNodes(titleDiv, findAnchor)

			if len(anchors) == 0 {
				article.Title = titleDiv.FirstChild.Data
				article.Link = ""
				continue
			}

			for _, anchor := range findNodes(titleDiv, findAnchor) {
				article.Title = anchor.FirstChild.Data
				link := pttHostURL + getAnchorLink(anchor)
				article.Link = link
				article.ID = article.ParseID(link)
			}
		}
		for _, metaDiv := range findNodes(articleBlock, findMetaDiv) {
			for _, date := range findNodes(metaDiv, findDateDiv) {
				article.Date = strings.TrimSpace(date.FirstChild.Data)
			}
			for _, author := range findNodes(metaDiv, findAuthorDiv) {
				article.Author = author.FirstChild.Data
			}
		}
		articles = append(articles, article)
		if isLastArticleBlock(articleBlock) {
			break
		}
	}
	return articles, nil
}

func isLastArticleBlock(articleBlock *html.Node) bool {
	for next := articleBlock.NextSibling; ; next = next.NextSibling {
		if next == nil {
			break
		}
		if next.Type == html.ElementNode {
			for _, attr := range next.Attr {
				if attr.Val == "r-list-sep" {
					return true
				}
				return false
			}
		}
	}
	return false
}

// FetchArticle build article object from html
func FetchArticle(board, articleCode string) (article.Article, error) {
	reqURL := makeArticleURL(board, articleCode)
	htmlNodes, err := fetchHTML(reqURL)
	if err != nil {
		return article.Article{}, err
	}
	atcl := article.Article{
		Link:  reqURL,
		Code:  articleCode,
		Board: board,
	}
	nodes := findNodes(htmlNodes, findOgTitleMeta)
	if len(nodes) > 0 {
		atcl.Title = getMetaContent(nodes[0])
	} else {
		atcl.Title = "[內文標題已被刪除]"
	}
	atcl.ID = atcl.ParseID(reqURL)
	pushBlocks := findNodes(htmlNodes, findPushBlocks)
	pushes := []article.Comment{}

PushBlocks:
	for _, pushBlock := range pushBlocks {
		push := article.Comment{}
		for _, pushIPDateTime := range findNodes(pushBlock, findPushIPDateTime) {
			ipdatetime := strings.TrimSpace(pushIPDateTime.FirstChild.Data)
			if ipdatetime == "" {
				continue PushBlocks
			}
			dateTime, err := parseDateTime(ipdatetime)
			if err != nil {
				log.WithFields(log.Fields{
					"ipdatetime": ipdatetime,
					"board":      board,
					"code":       articleCode,
				}).WithError(err).Warn("Parse DateTime Error")
				continue PushBlocks
			}
			push.DateTime = dateTime
		}
		for _, pushTag := range findNodes(pushBlock, findPushTag) {
			push.Tag = pushTag.FirstChild.Data
		}
		for _, pushUserID := range findNodes(pushBlock, findPushUserID) {
			push.UserID = pushUserID.FirstChild.Data
		}
		for _, pushContent := range findNodes(pushBlock, findPushContent) {
			content := pushContent.FirstChild.Data
			for n := pushContent.FirstChild.NextSibling; n != nil; n = n.NextSibling {
				if findEmailProtected(n) != nil {
					break
				}
				if n.FirstChild != nil {
					content += n.FirstChild.Data
				}
				if n.NextSibling != nil {
					content += n.NextSibling.Data
				}
			}
			push.Content = content
		}
		pushes = append(pushes, push)
	}
	if len(pushes) > 0 {
		atcl.LastPushDateTime = pushes[len(pushes)-1].DateTime
	}
	atcl.Comments = pushes
	return atcl, nil
}

func parseDateTime(ipdatetime string) (time.Time, error) {
	re, _ := regexp.Compile("(\\d+\\.\\d+\\.\\d+\\.\\d+)?\\s*(.*)")
	subMatches := re.FindStringSubmatch(ipdatetime)
	dateTime := strings.TrimSpace(subMatches[len(subMatches)-1])
	loc := time.FixedZone("CST", 8*60*60)
	t, err := time.ParseInLocation("01/02 15:04", dateTime, loc)
	if err != nil {
		return time.Time{}, err
	}
	t = t.AddDate(getYear(t), 0, 0)
	return t, nil
}

func getYear(commentTime time.Time) int {
	t := time.Now()
	commentTime = commentTime.AddDate(t.Year(), 0, 0)
	if commentTime.After(t) {
		return t.Year() - 1
	}
	return t.Year()
}

// CheckBoardExist use for checking board exist or not
func CheckBoardExist(board string) bool {
	return checkURLExist(makeBoardURL(board, -1))
}

// CheckArticleExist user for checking article exist or not
func CheckArticleExist(board, articleCode string) bool {
	return checkURLExist(makeArticleURL(board, articleCode))
}

func checkURLExist(url string) bool {
	resp, err := http.Get(url)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusOK {
		return true
	}
	return false
}

func makeBoardURL(board string, page int) string {
	var pageStr string
	if page < 0 {
		pageStr = ""
	} else {
		pageStr = strconv.Itoa(page)
	}
	return pttHostURL + "/bbs/" + board + "/index" + pageStr + ".html"
}

func makeArticleURL(board, articleCode string) string {
	return pttHostURL + "/bbs/" + board + "/" + articleCode + ".html"
}

// URLNotFoundError is an error type present 404 Not Found
type URLNotFoundError struct {
	URL string
}

func (u URLNotFoundError) Error() string {
	return "Fetched URL Not Found"
}

var errRedirect = errors.New("Redirect")

var client = &http.Client{
	CheckRedirect: func(req *http.Request, via []*http.Request) error {
		return errRedirect
	},
	Timeout: 30 * time.Second,
}

func fetchHTML(reqURL string) (doc *html.Node, err error) {
	resp, err := client.Get(reqURL)
	if err != nil && resp == nil {
		log.WithField("url", reqURL).WithError(err).Error("Fetch URL Failed")
		return nil, err
	}
	defer resp.Body.Close()

	if err == nil && resp.StatusCode == http.StatusNotFound {
		err = URLNotFoundError{reqURL}
		return nil, err
	}

	if uerr, ok := err.(*url.Error); ok && uerr.Err == errRedirect {
		req := passR18(reqURL)
		resp, err = client.Do(req)
		if err != nil {
			return nil, err
		}
	}

	doc, err = html.Parse(resp.Body)
	if err != nil {
		log.WithError(err).Error("Crawler Fetch HTML Failed")
		return nil, err
	}
	return doc, nil
}

func passR18(reqURL string) (req *http.Request) {

	req, _ = http.NewRequest("GET", reqURL, nil)

	over18Cookie := http.Cookie{
		Name:       "over18",
		Value:      "1",
		Domain:     "www.ptt.cc",
		Path:       "/",
		RawExpires: "Session",
		MaxAge:     0,
		HttpOnly:   false,
	}

	req.AddCookie(&over18Cookie)

	return req
}
<|endoftext|>


<commit_before>package entities

import (
	"encoding/json"
	"reflect"
	"testing"
	"time"

	"github.com/Vladimiroff/vec2d"
)

func TestMissionKey(t *testing.T) {
	startTime := time.Date(2012, time.November, 10, 23, 0, 0, 0, time.UTC).UnixNano() / 1e6
	mission := new(Mission)
	mission.Source = embeddedPlanet{planet.Name, planet.Owner, planet.Position}
	mission.StartTime = startTime

	if mission.Key() != "mission.1352588400000_GOP6720" {
		t.Error("Mission's key is ", mission.Key())
	}
}

func TestMissionMarshalling(t *testing.T) {
	var uMission Mission

	mMission, err := json.Marshal(mission)
	if err != nil {
		t.Error("Mission marshaling failed:", err)
	}

	err = json.Unmarshal(mMission, &uMission)
	if err != nil {
		t.Error("Mission unmarshaling failed:", err)
	}

	if mission.Key() != uMission.Key() {
		t.Error(
			"Keys of both missions are different!\n",
			mission.Key(),
			"!=",
			uMission.Key(),
		)
	}

	if !reflect.DeepEqual(mission, uMission) {
		t.Error("Missions are different after the marshal->unmarshal step")
	}
}

func TestEndAttackMission(t *testing.T) {
	var excessShips int32
	excessShips = secondMission.EndAttackMission(&endPlanet)
	if endPlanet.GetShipCount() != 12 {
		t.Error("End Planet ship count was expected  to be 12 but it is:", endPlanet.GetShipCount())
	}

	if endPlanet.Owner != "chochko" {
		t.Error("End Planet owner was expected  to be chochko but is:", endPlanet.Owner)
	}

	mission.ShipCount = 15
	excessShips = mission.EndAttackMission(&endPlanet)
	if endPlanet.GetShipCount() != 3 {
		t.Error("End Planet ship count was expected  to be 3 but it is:", endPlanet.GetShipCount())
	}

	if endPlanet.Owner != "gophie" {
		t.Error("End Planet owner was expected  to be gophie but is:", endPlanet.Owner)
	}

	if excessShips != 0 {
		t.Error("There shouldn't be any excess ships, but the value is", excessShips)
	}
}

func TestEndAttackMissionDenyTakeover(t *testing.T) {
	var excessShips int32
	endPlanet := new(Planet)
	*endPlanet = Planet{"", Color{0.59215686, 0.59215686, 0.59215686}, vec2d.New(2, 2), true, 6, 3, timeStamp, 2, 0, "chochko"}

	mission.ShipCount = 5
	excessShips = mission.EndAttackMission(endPlanet)
	if endPlanet.GetShipCount() != 0 {
		t.Error("End Planet ship count was expected to be 0 but it is:", endPlanet.GetShipCount())
	}

	if endPlanet.Owner != "chochko" {
		t.Error("End Planet owner was expected to be chochko but is:", endPlanet.Owner)
	}

	if excessShips != 5 {
		t.Error("There should be 5 excess ships, but the value is", excessShips)
	}
}

func TestTravelTime(t *testing.T) {
	source := vec2d.New(100, 200)
	target := vec2d.New(800, 150)
	expectedTime := time.Duration(7017)
	time := calculateTravelTime(source, target, 10)

	if time != expectedTime {
		t.Errorf(
			"CalculateTravelTime(%v, %v, 10) = %d instead of %d",
			source,
			target,
			time,
			expectedTime,
		)
	}
}
<commit_msg>Make the TestEndAttackMission leaderboard friendly<commit_after>package entities

import (
	"encoding/json"
	"reflect"
	"testing"
	"time"

	"github.com/Vladimiroff/vec2d"

	"warcluster/leaderboard"
)

func TestMissionKey(t *testing.T) {
	startTime := time.Date(2012, time.November, 10, 23, 0, 0, 0, time.UTC).UnixNano() / 1e6
	mission := new(Mission)
	mission.Source = embeddedPlanet{planet.Name, planet.Owner, planet.Position}
	mission.StartTime = startTime

	if mission.Key() != "mission.1352588400000_GOP6720" {
		t.Error("Mission's key is ", mission.Key())
	}
}

func TestMissionMarshalling(t *testing.T) {
	var uMission Mission

	mMission, err := json.Marshal(mission)
	if err != nil {
		t.Error("Mission marshaling failed:", err)
	}

	err = json.Unmarshal(mMission, &uMission)
	if err != nil {
		t.Error("Mission unmarshaling failed:", err)
	}

	if mission.Key() != uMission.Key() {
		t.Error(
			"Keys of both missions are different!\n",
			mission.Key(),
			"!=",
			uMission.Key(),
		)
	}

	if !reflect.DeepEqual(mission, uMission) {
		t.Error("Missions are different after the marshal->unmarshal step")
	}
}

func TestEndAttackMission(t *testing.T) {
	var excessShips int32

	leaderboard.Board = leaderboard.Leaderboard{
		{Username: "gophie", Planets: 10},
		{Username: "chochko", Planets: 8},
	}
	leaderboard.Places = map[string]int{
		"gophie":  0,
		"chochko": 1,
	}

	excessShips = secondMission.EndAttackMission(&endPlanet)
	if endPlanet.GetShipCount() != 12 {
		t.Error("End Planet ship count was expected  to be 12 but it is:", endPlanet.GetShipCount())
	}

	if endPlanet.Owner != "chochko" {
		t.Error("End Planet owner was expected  to be chochko but is:", endPlanet.Owner)
	}

	mission.ShipCount = 15
	excessShips = mission.EndAttackMission(&endPlanet)
	if endPlanet.GetShipCount() != 3 {
		t.Error("End Planet ship count was expected  to be 3 but it is:", endPlanet.GetShipCount())
	}

	if endPlanet.Owner != "gophie" {
		t.Error("End Planet owner was expected  to be gophie but is:", endPlanet.Owner)
	}

	if excessShips != 0 {
		t.Error("There shouldn't be any excess ships, but the value is", excessShips)
	}
}

func TestEndAttackMissionDenyTakeover(t *testing.T) {
	var excessShips int32
	endPlanet := new(Planet)
	*endPlanet = Planet{"", Color{0.59215686, 0.59215686, 0.59215686}, vec2d.New(2, 2), true, 6, 3, timeStamp, 2, 0, "chochko"}

	mission.ShipCount = 5
	excessShips = mission.EndAttackMission(endPlanet)
	if endPlanet.GetShipCount() != 0 {
		t.Error("End Planet ship count was expected to be 0 but it is:", endPlanet.GetShipCount())
	}

	if endPlanet.Owner != "chochko" {
		t.Error("End Planet owner was expected to be chochko but is:", endPlanet.Owner)
	}

	if excessShips != 5 {
		t.Error("There should be 5 excess ships, but the value is", excessShips)
	}
}

func TestTravelTime(t *testing.T) {
	source := vec2d.New(100, 200)
	target := vec2d.New(800, 150)
	expectedTime := time.Duration(7017)
	time := calculateTravelTime(source, target, 10)

	if time != expectedTime {
		t.Errorf(
			"CalculateTravelTime(%v, %v, 10) = %d instead of %d",
			source,
			target,
			time,
			expectedTime,
		)
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package bytes_test

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"sort"
	"unicode"
)

func ExampleBuffer() {
	var b bytes.Buffer // A Buffer needs no initialization.
	b.Write([]byte("Hello "))
	fmt.Fprintf(&b, "world!")
	b.WriteTo(os.Stdout)
	// Output: Hello world!
}

func ExampleBuffer_reader() {
	// A Buffer can turn a string or a []byte into an io.Reader.
	buf := bytes.NewBufferString("R29waGVycyBydWxlIQ==")
	dec := base64.NewDecoder(base64.StdEncoding, buf)
	io.Copy(os.Stdout, dec)
	// Output: Gophers rule!
}

func ExampleCompare() {
	// Interpret Compare's result by comparing it to zero.
	var a, b []byte
	if bytes.Compare(a, b) < 0 {
		// a less b
	}
	if bytes.Compare(a, b) <= 0 {
		// a less or equal b
	}
	if bytes.Compare(a, b) > 0 {
		// a greater b
	}
	if bytes.Compare(a, b) >= 0 {
		// a greater or equal b
	}

	// Prefer Equal to Compare for equality comparisons.
	if bytes.Equal(a, b) {
		// a equal b
	}
	if !bytes.Equal(a, b) {
		// a not equal b
	}
}

func ExampleCompare_search() {
	// Binary search to find a matching byte slice.
	var needle []byte
	var haystack [][]byte // Assume sorted
	i := sort.Search(len(haystack), func(i int) bool {
		// Return haystack[i] >= needle.
		return bytes.Compare(haystack[i], needle) >= 0
	})
	if i < len(haystack) && bytes.Equal(haystack[i], needle) {
		// Found it!
	}
}

func ExampleTrimSuffix() {
	var b = []byte("Hello, goodbye, etc!")
	b = bytes.TrimSuffix(b, []byte("goodbye, etc!"))
	b = bytes.TrimSuffix(b, []byte("gopher"))
	b = append(b, bytes.TrimSuffix([]byte("world!"), []byte("x!"))...)
	os.Stdout.Write(b)
	// Output: Hello, world!
}

func ExampleTrimPrefix() {
	var b = []byte("Goodbye,, world!")
	b = bytes.TrimPrefix(b, []byte("Goodbye,"))
	b = bytes.TrimPrefix(b, []byte("See ya,"))
	fmt.Printf("Hello%s", b)
	// Output: Hello, world!
}

func ExampleFields() {
	fmt.Printf("Fields are: %q", bytes.Fields([]byte("  foo bar  baz   ")))
	// Output: Fields are: ["foo" "bar" "baz"]
}

func ExampleFieldsFunc() {
	f := func(c rune) bool {
		return !unicode.IsLetter(c) && !unicode.IsNumber(c)
	}
	fmt.Printf("Fields are: %q", bytes.FieldsFunc([]byte("  foo1;bar2,baz3..."), f))
	// Output: Fields are: ["foo1" "bar2" "baz3"]
}

func ExampleContains() {
	fmt.Println(bytes.Contains([]byte("seafood"), []byte("foo")))
	fmt.Println(bytes.Contains([]byte("seafood"), []byte("bar")))
	fmt.Println(bytes.Contains([]byte("seafood"), []byte("")))
	fmt.Println(bytes.Contains([]byte(""), []byte("")))
	// Output:
	// true
	// false
	// true
	// true
}

func ExampleCount() {
	fmt.Println(bytes.Count([]byte("cheese"), []byte("e")))
	fmt.Println(bytes.Count([]byte("five"), []byte(""))) // before & after each rune
	// Output:
	// 3
	// 5
}

func ExampleEqualFold() {
	fmt.Println(bytes.EqualFold([]byte("Go"), []byte("go")))
	// Output: true
}

func ExampleHasPrefix() {
	fmt.Println(bytes.HasPrefix([]byte("Gopher"), []byte("Go")))
	fmt.Println(bytes.HasPrefix([]byte("Gopher"), []byte("C")))
	fmt.Println(bytes.HasPrefix([]byte("Gopher"), []byte("")))
	// Output:
	// true
	// false
	// true
}

func ExampleHasSuffix() {
	fmt.Println(bytes.HasSuffix([]byte("Amigo"), []byte("go")))
	fmt.Println(bytes.HasSuffix([]byte("Amigo"), []byte("O")))
	fmt.Println(bytes.HasSuffix([]byte("Amigo"), []byte("Ami")))
	fmt.Println(bytes.HasSuffix([]byte("Amigo"), []byte("")))
	// Output:
	// true
	// false
	// false
	// true
}

func ExampleIndex() {
	fmt.Println(bytes.Index([]byte("chicken"), []byte("ken")))
	fmt.Println(bytes.Index([]byte("chicken"), []byte("dmr")))
	// Output:
	// 4
	// -1
}

func ExampleIndexFunc() {
	f := func(c rune) bool {
		return unicode.Is(unicode.Han, c)
	}
	fmt.Println(bytes.IndexFunc([]byte("Hello, 世界"), f))
	fmt.Println(bytes.IndexFunc([]byte("Hello, world"), f))
	// Output:
	// 7
	// -1
}

func ExampleIndexAny() {
	fmt.Println(bytes.IndexAny([]byte("chicken"), "aeiouy"))
	fmt.Println(bytes.IndexAny([]byte("crwth"), "aeiouy"))
	// Output:
	// 2
	// -1
}

func ExampleIndexRune() {
	fmt.Println(bytes.IndexRune([]byte("chicken"), 'k'))
	fmt.Println(bytes.IndexRune([]byte("chicken"), 'd'))
	// Output:
	// 4
	// -1
}

func ExampleLastIndex() {
	fmt.Println(bytes.Index([]byte("go gopher"), []byte("go")))
	fmt.Println(bytes.LastIndex([]byte("go gopher"), []byte("go")))
	fmt.Println(bytes.LastIndex([]byte("go gopher"), []byte("rodent")))
	// Output:
	// 0
	// 3
	// -1
}

func ExampleJoin() {
	s := [][]byte{[]byte("foo"), []byte("bar"), []byte("baz")}
	fmt.Printf("%s", bytes.Join(s, []byte(", ")))
	// Output: foo, bar, baz
}

func ExampleRepeat() {
	fmt.Printf("ba%s", bytes.Repeat([]byte("na"), 2))
	// Output: banana
}

func ExampleReplace() {
	fmt.Printf("%s\n", bytes.Replace([]byte("oink oink oink"), []byte("k"), []byte("ky"), 2))
	fmt.Printf("%s\n", bytes.Replace([]byte("oink oink oink"), []byte("oink"), []byte("moo"), -1))
	// Output:
	// oinky oinky oink
	// moo moo moo
}

func ExampleSplit() {
	fmt.Printf("%q\n", bytes.Split([]byte("a,b,c"), []byte(",")))
	fmt.Printf("%q\n", bytes.Split([]byte("a man a plan a canal panama"), []byte("a ")))
	fmt.Printf("%q\n", bytes.Split([]byte(" xyz "), []byte("")))
	fmt.Printf("%q\n", bytes.Split([]byte(""), []byte("Bernardo O'Higgins")))
	// Output:
	// ["a" "b" "c"]
	// ["" "man " "plan " "canal panama"]
	// [" " "x" "y" "z" " "]
	// [""]
}

func ExampleSplitN() {
	fmt.Printf("%q\n", bytes.SplitN([]byte("a,b,c"), []byte(","), 2))
	z := bytes.SplitN([]byte("a,b,c"), []byte(","), 0)
	fmt.Printf("%q (nil = %v)\n", z, z == nil)
	// Output:
	// ["a" "b,c"]
	// [] (nil = true)
}

func ExampleSplitAfter() {
	fmt.Printf("%q\n", bytes.SplitAfter([]byte("a,b,c"), []byte(",")))
	// Output: ["a," "b," "c"]
}

func ExampleSplitAfterN() {
	fmt.Printf("%q\n", bytes.SplitAfterN([]byte("a,b,c"), []byte(","), 2))
	// Output: ["a," "b,c"]
}

func ExampleTitle() {
	fmt.Printf("%s", bytes.Title([]byte("her royal highness")))
	// Output: Her Royal Highness
}

func ExampleToTitle() {
	fmt.Printf("%s\n", bytes.ToTitle([]byte("loud noises")))
	fmt.Printf("%s\n", bytes.ToTitle([]byte("хлеб")))
	// Output:
	// LOUD NOISES
	// ХЛЕБ
}

func ExampleTrim() {
	fmt.Printf("[%q]", bytes.Trim([]byte(" !!! Achtung! Achtung! !!! "), "! "))
	// Output: ["Achtung! Achtung"]
}

func ExampleMap() {
	rot13 := func(r rune) rune {
		switch {
		case r >= 'A' && r <= 'Z':
			return 'A' + (r-'A'+13)%26
		case r >= 'a' && r <= 'z':
			return 'a' + (r-'a'+13)%26
		}
		return r
	}
	fmt.Printf("%s", bytes.Map(rot13, []byte("'Twas brillig and the slithy gopher...")))
	// Output: 'Gjnf oevyyvt naq gur fyvgul tbcure...
}

func ExampleTrimSpace() {
	fmt.Printf("%s", bytes.TrimSpace([]byte(" \t\n a lone gopher \n\t\r\n")))
	// Output: a lone gopher
}

func ExampleToUpper() {
	fmt.Printf("%s", bytes.ToUpper([]byte("Gopher")))
	// Output: GOPHER
}

func ExampleToLower() {
	fmt.Printf("%s", bytes.ToLower([]byte("Gopher")))
	// Output: gopher
}
<commit_msg>bytes: add example for (*Buffer).Grow<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package bytes_test

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"sort"
	"unicode"
)

func ExampleBuffer() {
	var b bytes.Buffer // A Buffer needs no initialization.
	b.Write([]byte("Hello "))
	fmt.Fprintf(&b, "world!")
	b.WriteTo(os.Stdout)
	// Output: Hello world!
}

func ExampleBuffer_reader() {
	// A Buffer can turn a string or a []byte into an io.Reader.
	buf := bytes.NewBufferString("R29waGVycyBydWxlIQ==")
	dec := base64.NewDecoder(base64.StdEncoding, buf)
	io.Copy(os.Stdout, dec)
	// Output: Gophers rule!
}

func ExampleBuffer_Grow() {
	var b bytes.Buffer
	b.Grow(64)
	bb := b.Bytes()
	b.Write([]byte("64 bytes or fewer"))
	fmt.Printf("%q", bb[:b.Len()])
	// Output: "64 bytes or fewer"
}

func ExampleCompare() {
	// Interpret Compare's result by comparing it to zero.
	var a, b []byte
	if bytes.Compare(a, b) < 0 {
		// a less b
	}
	if bytes.Compare(a, b) <= 0 {
		// a less or equal b
	}
	if bytes.Compare(a, b) > 0 {
		// a greater b
	}
	if bytes.Compare(a, b) >= 0 {
		// a greater or equal b
	}

	// Prefer Equal to Compare for equality comparisons.
	if bytes.Equal(a, b) {
		// a equal b
	}
	if !bytes.Equal(a, b) {
		// a not equal b
	}
}

func ExampleCompare_search() {
	// Binary search to find a matching byte slice.
	var needle []byte
	var haystack [][]byte // Assume sorted
	i := sort.Search(len(haystack), func(i int) bool {
		// Return haystack[i] >= needle.
		return bytes.Compare(haystack[i], needle) >= 0
	})
	if i < len(haystack) && bytes.Equal(haystack[i], needle) {
		// Found it!
	}
}

func ExampleTrimSuffix() {
	var b = []byte("Hello, goodbye, etc!")
	b = bytes.TrimSuffix(b, []byte("goodbye, etc!"))
	b = bytes.TrimSuffix(b, []byte("gopher"))
	b = append(b, bytes.TrimSuffix([]byte("world!"), []byte("x!"))...)
	os.Stdout.Write(b)
	// Output: Hello, world!
}

func ExampleTrimPrefix() {
	var b = []byte("Goodbye,, world!")
	b = bytes.TrimPrefix(b, []byte("Goodbye,"))
	b = bytes.TrimPrefix(b, []byte("See ya,"))
	fmt.Printf("Hello%s", b)
	// Output: Hello, world!
}

func ExampleFields() {
	fmt.Printf("Fields are: %q", bytes.Fields([]byte("  foo bar  baz   ")))
	// Output: Fields are: ["foo" "bar" "baz"]
}

func ExampleFieldsFunc() {
	f := func(c rune) bool {
		return !unicode.IsLetter(c) && !unicode.IsNumber(c)
	}
	fmt.Printf("Fields are: %q", bytes.FieldsFunc([]byte("  foo1;bar2,baz3..."), f))
	// Output: Fields are: ["foo1" "bar2" "baz3"]
}

func ExampleContains() {
	fmt.Println(bytes.Contains([]byte("seafood"), []byte("foo")))
	fmt.Println(bytes.Contains([]byte("seafood"), []byte("bar")))
	fmt.Println(bytes.Contains([]byte("seafood"), []byte("")))
	fmt.Println(bytes.Contains([]byte(""), []byte("")))
	// Output:
	// true
	// false
	// true
	// true
}

func ExampleCount() {
	fmt.Println(bytes.Count([]byte("cheese"), []byte("e")))
	fmt.Println(bytes.Count([]byte("five"), []byte(""))) // before & after each rune
	// Output:
	// 3
	// 5
}

func ExampleEqualFold() {
	fmt.Println(bytes.EqualFold([]byte("Go"), []byte("go")))
	// Output: true
}

func ExampleHasPrefix() {
	fmt.Println(bytes.HasPrefix([]byte("Gopher"), []byte("Go")))
	fmt.Println(bytes.HasPrefix([]byte("Gopher"), []byte("C")))
	fmt.Println(bytes.HasPrefix([]byte("Gopher"), []byte("")))
	// Output:
	// true
	// false
	// true
}

func ExampleHasSuffix() {
	fmt.Println(bytes.HasSuffix([]byte("Amigo"), []byte("go")))
	fmt.Println(bytes.HasSuffix([]byte("Amigo"), []byte("O")))
	fmt.Println(bytes.HasSuffix([]byte("Amigo"), []byte("Ami")))
	fmt.Println(bytes.HasSuffix([]byte("Amigo"), []byte("")))
	// Output:
	// true
	// false
	// false
	// true
}

func ExampleIndex() {
	fmt.Println(bytes.Index([]byte("chicken"), []byte("ken")))
	fmt.Println(bytes.Index([]byte("chicken"), []byte("dmr")))
	// Output:
	// 4
	// -1
}

func ExampleIndexFunc() {
	f := func(c rune) bool {
		return unicode.Is(unicode.Han, c)
	}
	fmt.Println(bytes.IndexFunc([]byte("Hello, 世界"), f))
	fmt.Println(bytes.IndexFunc([]byte("Hello, world"), f))
	// Output:
	// 7
	// -1
}

func ExampleIndexAny() {
	fmt.Println(bytes.IndexAny([]byte("chicken"), "aeiouy"))
	fmt.Println(bytes.IndexAny([]byte("crwth"), "aeiouy"))
	// Output:
	// 2
	// -1
}

func ExampleIndexRune() {
	fmt.Println(bytes.IndexRune([]byte("chicken"), 'k'))
	fmt.Println(bytes.IndexRune([]byte("chicken"), 'd'))
	// Output:
	// 4
	// -1
}

func ExampleLastIndex() {
	fmt.Println(bytes.Index([]byte("go gopher"), []byte("go")))
	fmt.Println(bytes.LastIndex([]byte("go gopher"), []byte("go")))
	fmt.Println(bytes.LastIndex([]byte("go gopher"), []byte("rodent")))
	// Output:
	// 0
	// 3
	// -1
}

func ExampleJoin() {
	s := [][]byte{[]byte("foo"), []byte("bar"), []byte("baz")}
	fmt.Printf("%s", bytes.Join(s, []byte(", ")))
	// Output: foo, bar, baz
}

func ExampleRepeat() {
	fmt.Printf("ba%s", bytes.Repeat([]byte("na"), 2))
	// Output: banana
}

func ExampleReplace() {
	fmt.Printf("%s\n", bytes.Replace([]byte("oink oink oink"), []byte("k"), []byte("ky"), 2))
	fmt.Printf("%s\n", bytes.Replace([]byte("oink oink oink"), []byte("oink"), []byte("moo"), -1))
	// Output:
	// oinky oinky oink
	// moo moo moo
}

func ExampleSplit() {
	fmt.Printf("%q\n", bytes.Split([]byte("a,b,c"), []byte(",")))
	fmt.Printf("%q\n", bytes.Split([]byte("a man a plan a canal panama"), []byte("a ")))
	fmt.Printf("%q\n", bytes.Split([]byte(" xyz "), []byte("")))
	fmt.Printf("%q\n", bytes.Split([]byte(""), []byte("Bernardo O'Higgins")))
	// Output:
	// ["a" "b" "c"]
	// ["" "man " "plan " "canal panama"]
	// [" " "x" "y" "z" " "]
	// [""]
}

func ExampleSplitN() {
	fmt.Printf("%q\n", bytes.SplitN([]byte("a,b,c"), []byte(","), 2))
	z := bytes.SplitN([]byte("a,b,c"), []byte(","), 0)
	fmt.Printf("%q (nil = %v)\n", z, z == nil)
	// Output:
	// ["a" "b,c"]
	// [] (nil = true)
}

func ExampleSplitAfter() {
	fmt.Printf("%q\n", bytes.SplitAfter([]byte("a,b,c"), []byte(",")))
	// Output: ["a," "b," "c"]
}

func ExampleSplitAfterN() {
	fmt.Printf("%q\n", bytes.SplitAfterN([]byte("a,b,c"), []byte(","), 2))
	// Output: ["a," "b,c"]
}

func ExampleTitle() {
	fmt.Printf("%s", bytes.Title([]byte("her royal highness")))
	// Output: Her Royal Highness
}

func ExampleToTitle() {
	fmt.Printf("%s\n", bytes.ToTitle([]byte("loud noises")))
	fmt.Printf("%s\n", bytes.ToTitle([]byte("хлеб")))
	// Output:
	// LOUD NOISES
	// ХЛЕБ
}

func ExampleTrim() {
	fmt.Printf("[%q]", bytes.Trim([]byte(" !!! Achtung! Achtung! !!! "), "! "))
	// Output: ["Achtung! Achtung"]
}

func ExampleMap() {
	rot13 := func(r rune) rune {
		switch {
		case r >= 'A' && r <= 'Z':
			return 'A' + (r-'A'+13)%26
		case r >= 'a' && r <= 'z':
			return 'a' + (r-'a'+13)%26
		}
		return r
	}
	fmt.Printf("%s", bytes.Map(rot13, []byte("'Twas brillig and the slithy gopher...")))
	// Output: 'Gjnf oevyyvt naq gur fyvgul tbcure...
}

func ExampleTrimSpace() {
	fmt.Printf("%s", bytes.TrimSpace([]byte(" \t\n a lone gopher \n\t\r\n")))
	// Output: a lone gopher
}

func ExampleToUpper() {
	fmt.Printf("%s", bytes.ToUpper([]byte("Gopher")))
	// Output: GOPHER
}

func ExampleToLower() {
	fmt.Printf("%s", bytes.ToLower([]byte("Gopher")))
	// Output: gopher
}
<|endoftext|>


<commit_before>// +build windows

package jobsupervisor

import (
	"os"

	boshhandler "github.com/cloudfoundry/bosh-agent/handler"
	boshmonit "github.com/cloudfoundry/bosh-agent/jobsupervisor/monit"
	boshplatform "github.com/cloudfoundry/bosh-agent/platform"
	boshdir "github.com/cloudfoundry/bosh-agent/settings/directories"
	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
)

const jobSupervisorListenPort = 2825

type Provider struct {
	supervisors map[string]JobSupervisor
}

func NewProvider(
	platform boshplatform.Platform,
	client boshmonit.Client,
	logger boshlog.Logger,
	dirProvider boshdir.Provider,
	handler boshhandler.Handler,
) (p Provider) {
	fs := platform.GetFs()
	runner := platform.GetRunner()

	network, err := platform.GetDefaultNetwork()
	var machineIP string
	if err != nil {
		machineIP, _ = os.Hostname()
		logger.Debug("providerWindows", "Initializing jobsupervisor.provider_windows: %s, using hostname \"%s\"instead of IP", err, machineIP)
	} else {
		machineIP = network.IP
	}

	p.supervisors = map[string]JobSupervisor{
		"dummy":      NewDummyJobSupervisor(),
		"dummy-nats": NewDummyNatsJobSupervisor(handler),
		"windows":    NewWrapperJobSupervisor(NewWindowsJobSupervisor(runner, dirProvider, fs, logger, jobSupervisorListenPort, make(chan bool), machineIP), fs, dirProvider, logger),
	}

	return
}

func (p Provider) Get(name string) (supervisor JobSupervisor, err error) {
	supervisor, found := p.supervisors[name]
	if !found {
		err = bosherr.Errorf("JobSupervisor %s could not be found", name)
	}
	return
}
<commit_msg>Backward compatible monit support for windows job provider.<commit_after>// +build windows

package jobsupervisor

import (
	"os"

	boshhandler "github.com/cloudfoundry/bosh-agent/handler"
	boshmonit "github.com/cloudfoundry/bosh-agent/jobsupervisor/monit"
	boshplatform "github.com/cloudfoundry/bosh-agent/platform"
	boshdir "github.com/cloudfoundry/bosh-agent/settings/directories"
	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
)

const jobSupervisorListenPort = 2825

type Provider struct {
	supervisors map[string]JobSupervisor
}

func NewProvider(
	platform boshplatform.Platform,
	client boshmonit.Client,
	logger boshlog.Logger,
	dirProvider boshdir.Provider,
	handler boshhandler.Handler,
) (p Provider) {
	fs := platform.GetFs()
	runner := platform.GetRunner()

	network, err := platform.GetDefaultNetwork()
	var machineIP string
	if err != nil {
		machineIP, _ = os.Hostname()
		logger.Debug("providerWindows", "Initializing jobsupervisor.provider_windows: %s, using hostname \"%s\"instead of IP", err, machineIP)
	} else {
		machineIP = network.IP
	}

	p.supervisors = map[string]JobSupervisor{
		"monit":    NewWrapperJobSupervisor(NewWindowsJobSupervisor(runner, dirProvider, fs, logger, jobSupervisorListenPort, make(chan bool), machineIP), fs, dirProvider, logger),
		"dummy":      NewDummyJobSupervisor(),
		"dummy-nats": NewDummyNatsJobSupervisor(handler),
		"windows":    NewWrapperJobSupervisor(NewWindowsJobSupervisor(runner, dirProvider, fs, logger, jobSupervisorListenPort, make(chan bool), machineIP), fs, dirProvider, logger),
	}

	return
}

func (p Provider) Get(name string) (supervisor JobSupervisor, err error) {
	supervisor, found := p.supervisors[name]
	if !found {
		err = bosherr.Errorf("JobSupervisor %s could not be found", name)
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/geek1011/koboutils/kobo"

	"github.com/beevik/etree"
	"github.com/mattn/go-zglob"
	"golang.org/x/tools/godoc/vfs/zipfs"

	"github.com/spf13/pflag"

	_ "github.com/mattn/go-sqlite3"
)

var version = "dev"

func helpExit() {
	fmt.Fprintf(os.Stderr, "Usage: seriesmeta [OPTIONS] [KOBO_PATH]\n\nVersion:\n  seriesmeta %s\n\nOptions:\n", version)
	pflag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\nArguments:\n  KOBO_PATH is the path to the Kobo eReader. If not specified, seriesmeta will try to automatically detect the Kobo.\n")
	if runtime.GOOS == "windows" {
		time.Sleep(time.Second * 2)
	}
	os.Exit(1)
}

func errExit() {
	if runtime.GOOS == "windows" {
		time.Sleep(time.Second * 2)
	}
	os.Exit(1)
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

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

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}
	return out.Close()
}

// pathToContentID gets the content ID for a book. The path needs to be relative to the root of the kobo.
func pathToContentID(relpath string) string {
	return fmt.Sprintf("file:///mnt/onboard/%s", filepath.ToSlash(relpath))
}

func contentIDToImageID(contentID string) string {
	imageID := contentID

	imageID = strings.Replace(imageID, " ", "_", -1)
	imageID = strings.Replace(imageID, "/", "_", -1)
	imageID = strings.Replace(imageID, ":", "_", -1)
	imageID = strings.Replace(imageID, ".", "_", -1)

	return imageID
}

func getMeta(path string) (string, float64, error) {
	zr, err := zip.OpenReader(path)
	if err != nil {
		return "", 0, err
	}

	zfs := zipfs.New(zr, "epub")
	rsk, err := zfs.Open("/META-INF/container.xml")
	if err != nil {
		return "", 0, err
	}
	defer rsk.Close()

	container := etree.NewDocument()
	_, err = container.ReadFrom(rsk)
	if err != nil {
		return "", 0, err
	}

	rootfile := ""
	for _, e := range container.FindElements("//rootfiles/rootfile[@full-path]") {
		rootfile = e.SelectAttrValue("full-path", "")
	}

	if rootfile == "" {
		return "", 0, errors.New("Cannot parse container")
	}

	rrsk, err := zfs.Open("/" + rootfile)
	if err != nil {
		return "", 0, err
	}
	defer rrsk.Close()

	opf := etree.NewDocument()
	_, err = opf.ReadFrom(rrsk)
	if err != nil {
		return "", 0, err
	}

	var series string
	for _, e := range opf.FindElements("//meta[@name='calibre:series']") {
		series = e.SelectAttrValue("content", "")
		break
	}

	var seriesNumber float64
	for _, e := range opf.FindElements("//meta[@name='calibre:series_index']") {
		i, err := strconv.ParseFloat(e.SelectAttrValue("content", "0"), 64)
		if err == nil {
			seriesNumber = i
			break
		}
	}

	return series, seriesNumber, nil
}

func main() {
	help := pflag.BoolP("help", "h", false, "Show this help message")
	pflag.Parse()

	if *help || pflag.NArg() > 1 {
		helpExit()
	}

	log := func(format string, a ...interface{}) {
		fmt.Printf(format, a...)
	}

	logE := func(format string, a ...interface{}) {
		fmt.Fprintf(os.Stderr, format, a...)
	}

	var kpath string
	if pflag.NArg() == 1 {
		kpath = strings.Replace(pflag.Arg(0), ".kobo", "", 1)
	} else {
		log("No kobo specified, attempting to detect one\n")
		kobos, err := kobo.Find()
		if err != nil {
			logE("Fatal: could not automatically detect a kobo: %v\n", err)
			errExit()
		} else if len(kobos) < 1 {
			logE("Fatal: could not automatically detect a kobo\n")
			errExit()
		}
		kpath = kobos[0]
	}

	log("Checking kobo at '%s'\n", kpath)
	if !kobo.IsKobo(kpath) {
		logE("Fatal: '%s' is not a valid kobo\n", kpath)
	}

	kpath, err := filepath.Abs(kpath)
	if err != nil {
		logE("Fatal: Could not resolve path to kobo\n")
		errExit()
	}

	dbpath := filepath.Join(kpath, ".kobo", "KoboReader.sqlite")

	log("Making backup of KoboReader.sqlite\n")
	err = copyFile(dbpath, dbpath+".bak")
	if err != nil {
		logE("Fatal: Could not make copy of KoboReader.sqlite: %v\n", err)
		errExit()
	}

	log("Opening KoboReader.sqlite\n")
	db, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		logE("Fatal: Could not open KoboReader.sqlite: %v\n", err)
		errExit()
	}

	log("Searching for sideloaded epubs and kepubs\n")
	epubs, err := zglob.Glob(filepath.Join(kpath, "**", "*.epub"))
	if err != nil {
		logE("Fatal: Could not search for epubs: %v\n", err)
		errExit()
	}

	log("\nUpdating metadata for %d books\n", len(epubs))
	errcount := 0
	for i, epub := range epubs {
		rpath, err := filepath.Rel(kpath, epub)
		if err != nil {
			log("[%d/%d] Updating '%s'\n", i+1, len(epubs), epub)
			logE("  Error: could not resolve path: %v\n", err)
			errcount++
			continue
		}

		log("[%d/%d] Updating '%s'\n", i+1, len(epubs), rpath)
		series, seriesNumber, err := getMeta(epub)
		if err != nil {
			logE("  Error: could not read metadata: %v\n", err)
			errcount++
			continue
		}

		if series == "" && seriesNumber == 0 {
			log("  No series\n")
			continue
		}

		log("  Series: '%s' Number: %v\n", series, seriesNumber)

		iid := contentIDToImageID(pathToContentID(rpath))

		res, err := db.Exec("UPDATE content SET Series=?, SeriesNumber=? WHERE ImageID=?", sql.NullString{
			String: series,
			Valid:  series != "",
		}, sql.NullString{
			String: fmt.Sprintf("%v", seriesNumber),
			Valid:  seriesNumber > 0,
		}, iid)
		if err != nil {
			logE("  Error: could not update database: %v\n", err)
			errcount++
			continue
		}

		ra, err := res.RowsAffected()
		if err != nil {
			logE("  Error: could not update database: %v\n", err)
			errcount++
			continue
		}

		if ra > 1 {
			logE("  Warn: more than one match in database for ImageID\n")
		} else if ra < 1 {
			logE("  Error: could not update database: no entry in database for book (the kobo may still need to import the book)\n")
			errcount++
			continue
		}
	}

	time.Sleep(time.Second)
	log("\nFinished updating metadata. %v books processed. %v errors.\n", len(epubs), errcount)
}
<commit_msg>Improved seriesmeta output format<commit_after>package main

import (
	"archive/zip"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/geek1011/koboutils/kobo"

	"github.com/beevik/etree"
	"github.com/mattn/go-zglob"
	"golang.org/x/tools/godoc/vfs/zipfs"

	"github.com/spf13/pflag"

	_ "github.com/mattn/go-sqlite3"
)

var version = "dev"

func helpExit() {
	fmt.Fprintf(os.Stderr, "Usage: seriesmeta [OPTIONS] [KOBO_PATH]\n\nVersion:\n  seriesmeta %s\n\nOptions:\n", version)
	pflag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\nArguments:\n  KOBO_PATH is the path to the Kobo eReader. If not specified, seriesmeta will try to automatically detect the Kobo.\n")
	if runtime.GOOS == "windows" {
		time.Sleep(time.Second * 2)
	}
	os.Exit(1)
}

func errExit() {
	if runtime.GOOS == "windows" {
		time.Sleep(time.Second * 2)
	}
	os.Exit(1)
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

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

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}
	return out.Close()
}

// pathToContentID gets the content ID for a book. The path needs to be relative to the root of the kobo.
func pathToContentID(relpath string) string {
	return fmt.Sprintf("file:///mnt/onboard/%s", filepath.ToSlash(relpath))
}

func contentIDToImageID(contentID string) string {
	imageID := contentID

	imageID = strings.Replace(imageID, " ", "_", -1)
	imageID = strings.Replace(imageID, "/", "_", -1)
	imageID = strings.Replace(imageID, ":", "_", -1)
	imageID = strings.Replace(imageID, ".", "_", -1)

	return imageID
}

func getMeta(path string) (string, float64, error) {
	zr, err := zip.OpenReader(path)
	if err != nil {
		return "", 0, err
	}

	zfs := zipfs.New(zr, "epub")
	rsk, err := zfs.Open("/META-INF/container.xml")
	if err != nil {
		return "", 0, err
	}
	defer rsk.Close()

	container := etree.NewDocument()
	_, err = container.ReadFrom(rsk)
	if err != nil {
		return "", 0, err
	}

	rootfile := ""
	for _, e := range container.FindElements("//rootfiles/rootfile[@full-path]") {
		rootfile = e.SelectAttrValue("full-path", "")
	}

	if rootfile == "" {
		return "", 0, errors.New("Cannot parse container")
	}

	rrsk, err := zfs.Open("/" + rootfile)
	if err != nil {
		return "", 0, err
	}
	defer rrsk.Close()

	opf := etree.NewDocument()
	_, err = opf.ReadFrom(rrsk)
	if err != nil {
		return "", 0, err
	}

	var series string
	for _, e := range opf.FindElements("//meta[@name='calibre:series']") {
		series = e.SelectAttrValue("content", "")
		break
	}

	var seriesNumber float64
	for _, e := range opf.FindElements("//meta[@name='calibre:series_index']") {
		i, err := strconv.ParseFloat(e.SelectAttrValue("content", "0"), 64)
		if err == nil {
			seriesNumber = i
			break
		}
	}

	return series, seriesNumber, nil
}

func main() {
	help := pflag.BoolP("help", "h", false, "Show this help message")
	pflag.Parse()

	if *help || pflag.NArg() > 1 {
		helpExit()
	}

	log := func(format string, a ...interface{}) {
		fmt.Printf(format, a...)
	}

	logE := func(format string, a ...interface{}) {
		fmt.Fprintf(os.Stderr, format, a...)
	}

	var kpath string
	if pflag.NArg() == 1 {
		kpath = strings.Replace(pflag.Arg(0), ".kobo", "", 1)
	} else {
		log("No kobo specified, attempting to detect one\n")
		kobos, err := kobo.Find()
		if err != nil {
			logE("Fatal: could not automatically detect a kobo: %v\n", err)
			errExit()
		} else if len(kobos) < 1 {
			logE("Fatal: could not automatically detect a kobo\n")
			errExit()
		}
		kpath = kobos[0]
	}

	log("Checking kobo at '%s'\n", kpath)
	if !kobo.IsKobo(kpath) {
		logE("Fatal: '%s' is not a valid kobo\n", kpath)
	}

	kpath, err := filepath.Abs(kpath)
	if err != nil {
		logE("Fatal: Could not resolve path to kobo\n")
		errExit()
	}

	dbpath := filepath.Join(kpath, ".kobo", "KoboReader.sqlite")

	log("Making backup of KoboReader.sqlite\n")
	err = copyFile(dbpath, dbpath+".bak")
	if err != nil {
		logE("Fatal: Could not make copy of KoboReader.sqlite: %v\n", err)
		errExit()
	}

	log("Opening KoboReader.sqlite\n")
	db, err := sql.Open("sqlite3", dbpath)
	if err != nil {
		logE("Fatal: Could not open KoboReader.sqlite: %v\n", err)
		errExit()
	}

	log("Searching for sideloaded epubs and kepubs\n")
	epubs, err := zglob.Glob(filepath.Join(kpath, "**", "*.epub"))
	if err != nil {
		logE("Fatal: Could not search for epubs: %v\n", err)
		errExit()
	}

	log("\nUpdating metadata for %d books\n", len(epubs))
	var updated, nometa, errcount int
	digits := len(fmt.Sprint(len(epubs)))
	numFmt, spFmt := fmt.Sprintf("[%%%dd/%d] ", digits, len(epubs)), strings.Repeat(" ", (digits*2)+4)
	for i, epub := range epubs {
		rpath, err := filepath.Rel(kpath, epub)
		if err != nil {
			log(numFmt+"%s\n", i+1, epub)
			logE(spFmt+"Error: could not resolve path: %v\n", err)
			errcount++
			continue
		}

		log(numFmt+"%s\n", i+1, rpath)
		series, seriesNumber, err := getMeta(epub)
		if err != nil {
			logE(spFmt+"Error: could not read metadata: %v\n", err)
			errcount++
			continue
		}

		if series == "" && seriesNumber == 0 {
			nometa++
			continue
		}

		log(spFmt+"(%s, %v)\n", series, seriesNumber)

		iid := contentIDToImageID(pathToContentID(rpath))

		res, err := db.Exec("UPDATE content SET Series=?, SeriesNumber=? WHERE ImageID=?", sql.NullString{
			String: series,
			Valid:  series != "",
		}, sql.NullString{
			String: fmt.Sprintf("%v", seriesNumber),
			Valid:  seriesNumber > 0,
		}, iid)
		if err != nil {
			logE(spFmt+"Error: could not update database: %v\n", err)
			errcount++
			continue
		}

		ra, err := res.RowsAffected()
		if err != nil {
			logE(spFmt+"Error: could not update database: %v\n", err)
			errcount++
			continue
		}

		if ra > 1 {
			logE(spFmt + "Warn: more than one match in database for ImageID\n")
		} else if ra < 1 {
			logE(spFmt + "Error: could not update database: no entry in database for book (the kobo may still need to import the book)\n")
			errcount++
			continue
		}

		updated++
	}

	time.Sleep(time.Second)
	log("\nFinished updating metadata. %d updated, %d without metadata, %d errored.\n", updated, nometa, errcount)
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/tools/godoc/vfs/zipfs"

	"github.com/beevik/etree"
	_ "github.com/mattn/go-sqlite3"
)

func pathToContentID(koboPath, path string) (imageID string, err error) {
	relPath, err := filepath.Rel(koboPath, path)
	if err != nil {
		return "", fmt.Errorf("could not get relative path to file: %v", err)
	}

	contentID := fmt.Sprintf("file:///mnt/onboard/%s", relPath)

	return contentID, nil
}

func contentIDToImageID(contentID string) string {
	imageID := contentID

	imageID = strings.Replace(imageID, " ", "_", -1)
	imageID = strings.Replace(imageID, "/", "_", -1)
	imageID = strings.Replace(imageID, ":", "_", -1)
	imageID = strings.Replace(imageID, ".", "_", -1)

	return imageID
}

func updateSeriesMeta(db *sql.DB, imageID, series string, seriesNumber float64) (int64, error) {
	res, err := db.Exec("UPDATE content SET Series=?, SeriesNumber=? WHERE ImageID=?", sql.NullString{
		String: series,
		Valid:  series != "",
	}, sql.NullString{
		String: fmt.Sprintf("%v", seriesNumber),
		Valid:  seriesNumber > 0,
	}, imageID)

	if err != nil {
		return 0, err
	}

	return res.RowsAffected()
}

func getEPUBMeta(path string) (string, float64, error) {
	zr, err := zip.OpenReader(path)
	if err != nil {
		return "", 0, err
	}

	zfs := zipfs.New(zr, "epub")
	rsk, err := zfs.Open("/META-INF/container.xml")
	if err != nil {
		return "", 0, err
	}
	defer rsk.Close()

	container := etree.NewDocument()
	_, err = container.ReadFrom(rsk)
	if err != nil {
		return "", 0, err
	}

	rootfile := ""
	for _, e := range container.FindElements("//rootfiles/rootfile[@full-path]") {
		rootfile = e.SelectAttrValue("full-path", "")
	}

	if rootfile == "" {
		return "", 0, errors.New("Cannot parse container")
	}

	rrsk, err := zfs.Open("/" + rootfile)
	if err != nil {
		return "", 0, err
	}
	defer rrsk.Close()

	opf := etree.NewDocument()
	_, err = opf.ReadFrom(rrsk)
	if err != nil {
		return "", 0, err
	}

	var series string
	for _, e := range opf.FindElements("//meta[@name='calibre:series']") {
		series = e.SelectAttrValue("content", "")
		break
	}

	var seriesNumber float64
	for _, e := range opf.FindElements("//meta[@name='calibre:series_index']") {
		i, err := strconv.ParseFloat(e.SelectAttrValue("content", "0"), 64)
		if err == nil {
			seriesNumber = i
			break
		}
	}

	return series, seriesNumber, nil
}

func updateSeriesMetaFromEPUB(db *sql.DB, koboPath, epubPath string) (int64, error) {
	series, seriesNumber, err := getEPUBMeta(epubPath)
	if err != nil {
		return 0, err
	}

	cid, err := pathToContentID(koboPath, epubPath)
	if err != nil {
		return 0, err
	}

	iid := contentIDToImageID(cid)

	fmt.Printf("INFO: UPDATE %s => [%s %v]\n", iid, series, seriesNumber)

	return updateSeriesMeta(db, iid, series, seriesNumber)
}

func loadKoboDB(koboPath string) (*sql.DB, error) {
	koboDBPath := filepath.Join(koboPath, ".kobo/KoboReader.sqlite")

	if _, err := os.Stat(koboDBPath); os.IsNotExist(err) {
		return nil, fmt.Errorf("Kobo database %s does not exist", koboDBPath)
	}

	return sql.Open("sqlite3", koboDBPath)
}

func main() {
	if len(os.Args) != 3 {
		fmt.Printf("USAGE: %s KOBO_ROOT_PATH EPUB_PATH\n", filepath.Base(os.Args[0]))
		os.Exit(1)
	}

	koboPath, err := filepath.Abs(os.Args[1])
	if err != nil {
		fmt.Printf("FATAL: Could resolve Kobo path %s: %v\n", os.Args[1], err)
		os.Exit(1)
	}

	epubPath, err := filepath.Abs(os.Args[2])
	if err != nil {
		fmt.Printf("FATAL: Could resolve ePub path %s: %v\n", os.Args[2], err)
		os.Exit(1)
	}

	if !strings.HasPrefix(epubPath, koboPath) {
		fmt.Printf("FATAL: ePub file not in the specified Kobo path.\n")
		os.Exit(1)
	}

	db, err := loadKoboDB(koboPath)
	if err != nil {
		fmt.Printf("FATAL: Could not open Kobo database: %v\n", err)
		os.Exit(1)
	}

	ra, err := updateSeriesMetaFromEPUB(db, koboPath, epubPath)
	if err != nil {
		fmt.Printf("ERROR: Could not update series metadata from epub: %v\n", err)
	} else {
		if ra < 1 {
			fmt.Printf("ERROR: Could not update series metadata from epub: the database does not yet have an entry for the specified book. Please let the kobo import the book before using this tool.\n")
		}
		if ra > 1 {
			fmt.Printf("WARN: More than 1 match for book in database.\n")
		}
	}
}
<commit_msg>Added support for updating series metadata for all books<commit_after>package main

import (
	"archive/zip"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/mattn/go-zglob"

	"golang.org/x/tools/godoc/vfs/zipfs"

	"github.com/beevik/etree"
	_ "github.com/mattn/go-sqlite3"
)

func pathToContentID(koboPath, path string) (imageID string, err error) {
	relPath, err := filepath.Rel(koboPath, path)
	if err != nil {
		return "", fmt.Errorf("could not get relative path to file: %v", err)
	}

	contentID := fmt.Sprintf("file:///mnt/onboard/%s", relPath)

	return contentID, nil
}

func contentIDToImageID(contentID string) string {
	imageID := contentID

	imageID = strings.Replace(imageID, " ", "_", -1)
	imageID = strings.Replace(imageID, "/", "_", -1)
	imageID = strings.Replace(imageID, ":", "_", -1)
	imageID = strings.Replace(imageID, ".", "_", -1)

	return imageID
}

func updateSeriesMeta(db *sql.DB, imageID, series string, seriesNumber float64) (int64, error) {
	res, err := db.Exec("UPDATE content SET Series=?, SeriesNumber=? WHERE ImageID=?", sql.NullString{
		String: series,
		Valid:  series != "",
	}, sql.NullString{
		String: fmt.Sprintf("%v", seriesNumber),
		Valid:  seriesNumber > 0,
	}, imageID)

	if err != nil {
		return 0, err
	}

	return res.RowsAffected()
}

func getEPUBMeta(path string) (string, float64, error) {
	zr, err := zip.OpenReader(path)
	if err != nil {
		return "", 0, err
	}

	zfs := zipfs.New(zr, "epub")
	rsk, err := zfs.Open("/META-INF/container.xml")
	if err != nil {
		return "", 0, err
	}
	defer rsk.Close()

	container := etree.NewDocument()
	_, err = container.ReadFrom(rsk)
	if err != nil {
		return "", 0, err
	}

	rootfile := ""
	for _, e := range container.FindElements("//rootfiles/rootfile[@full-path]") {
		rootfile = e.SelectAttrValue("full-path", "")
	}

	if rootfile == "" {
		return "", 0, errors.New("Cannot parse container")
	}

	rrsk, err := zfs.Open("/" + rootfile)
	if err != nil {
		return "", 0, err
	}
	defer rrsk.Close()

	opf := etree.NewDocument()
	_, err = opf.ReadFrom(rrsk)
	if err != nil {
		return "", 0, err
	}

	var series string
	for _, e := range opf.FindElements("//meta[@name='calibre:series']") {
		series = e.SelectAttrValue("content", "")
		break
	}

	var seriesNumber float64
	for _, e := range opf.FindElements("//meta[@name='calibre:series_index']") {
		i, err := strconv.ParseFloat(e.SelectAttrValue("content", "0"), 64)
		if err == nil {
			seriesNumber = i
			break
		}
	}

	return series, seriesNumber, nil
}

func updateSeriesMetaFromEPUB(db *sql.DB, koboPath, epubPath string) (int64, error) {
	series, seriesNumber, err := getEPUBMeta(epubPath)
	if err != nil {
		return 0, err
	}

	cid, err := pathToContentID(koboPath, epubPath)
	if err != nil {
		return 0, err
	}

	iid := contentIDToImageID(cid)

	fmt.Printf("INFO: UPDATE %s => [%s %v]\n", iid, series, seriesNumber)

	return updateSeriesMeta(db, iid, series, seriesNumber)
}

func loadKoboDB(koboPath string) (*sql.DB, error) {
	koboDBPath := filepath.Join(koboPath, ".kobo/KoboReader.sqlite")

	if _, err := os.Stat(koboDBPath); os.IsNotExist(err) {
		return nil, fmt.Errorf("Kobo database %s does not exist", koboDBPath)
	}

	return sql.Open("sqlite3", koboDBPath)
}

func main() {
	if len(os.Args) < 2 || len(os.Args) > 3 {
		fmt.Printf("USAGE: %s KOBO_ROOT_PATH [EPUB_PATH]\n", filepath.Base(os.Args[0]))
		os.Exit(1)
	}

	koboPath, err := filepath.Abs(os.Args[1])
	if err != nil {
		fmt.Printf("FATAL: Could resolve Kobo path %s: %v\n", os.Args[1], err)
		os.Exit(1)
	}

	if _, err := os.Stat(filepath.Join(koboPath, ".kobo")); os.IsNotExist(err) {
		fmt.Printf("FATAL: %s is not a valid path to a Kobo eReader.\n", os.Args[1])
		fmt.Printf("USAGE: %s KOBO_ROOT_PATH [EPUB_PATH]\n", filepath.Base(os.Args[0]))
		os.Exit(1)
	}

	if len(os.Args) == 3 {
		epubPath, err := filepath.Abs(os.Args[2])
		if err != nil {
			fmt.Printf("FATAL: Could resolve ePub path %s: %v\n", os.Args[2], err)
			os.Exit(1)
		}

		if !strings.HasPrefix(epubPath, koboPath) {
			fmt.Printf("FATAL: ePub file not in the specified Kobo path.\n")
			os.Exit(1)
		}

		db, err := loadKoboDB(koboPath)
		if err != nil {
			fmt.Printf("FATAL: Could not open Kobo database: %v\n", err)
			os.Exit(1)
		}

		ra, err := updateSeriesMetaFromEPUB(db, koboPath, epubPath)
		if err != nil {
			fmt.Printf("ERROR: Could not update series metadata from epub: %v\n", err)
			os.Exit(1)
		} else if ra < 1 {
			fmt.Printf("ERROR: Could not update series metadata from epub: the database does not yet have an entry for the specified book. Please let the kobo import the book before using this tool.\n")
		} else if ra > 1 {
			fmt.Printf("WARN: More than 1 match for book in database.\n")
		}
	} else {
		db, err := loadKoboDB(koboPath)
		if err != nil {
			fmt.Printf("FATAL: Could not open Kobo database: %v\n", err)
			os.Exit(1)
		}

		matches, err := zglob.Glob(filepath.Join(koboPath, "**/*.epub"))
		if err != nil {
			fmt.Printf("FATAL: Error searching for epub files: %v\n", err)
			os.Exit(1)
		}

		epubs := []string{}
		for _, match := range matches {
			if strings.HasPrefix(filepath.Base(match), ".") {
				continue
			}
			epubs = append(epubs, match)
		}

		fmt.Printf("INFO: Found %v epub files", len(epubs))

		errcount := 0
		for _, epub := range epubs {
			ra, err := updateSeriesMetaFromEPUB(db, koboPath, epub)
			if err != nil {
				fmt.Printf("ERROR: Could not update series metadata from epub: %v\n", err)
				errcount++
			} else if ra < 1 {
				fmt.Printf("ERROR: Could not update series metadata from epub: the database does not yet have an entry for the specified book. Please let the kobo import the book before using this tool.\n")
				errcount++
			} else if ra > 1 {
				fmt.Printf("WARN: More than 1 match for book in database.\n")
			}
			fmt.Println()
		}

		fmt.Printf("INFO: Finished updating metadata. %v books processed. %v errors.", len(epubs), errcount)
	}
}
<|endoftext|>


<commit_before>package server

// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Meh, we need to get rid of this shit

import (
	"errors"
	"io"
	"log"
	"reflect"
	"strings"
	"sync"
	"unicode"
	"unicode/utf8"

	"golang.org/x/net/context"
)

var (
	lastStreamResponseError = errors.New("EOS")
	// A value sent as a placeholder for the server's response value when the server
	// receives an invalid request. It is never decoded by the client since the Response
	// contains an error when it is used.
	invalidRequest = struct{}{}

	// Precompute the reflect type for error. Can't use error directly
	// because Typeof takes an empty interface value. This is annoying.
	typeOfError = reflect.TypeOf((*error)(nil)).Elem()
)

type methodType struct {
	sync.Mutex  // protects counters
	method      reflect.Method
	ArgType     reflect.Type
	ReplyType   reflect.Type
	ContextType reflect.Type
	stream      bool
	numCalls    uint
}

func (m *methodType) NumCalls() (n uint) {
	m.Lock()
	n = m.numCalls
	m.Unlock()
	return n
}

type service struct {
	name   string                 // name of service
	rcvr   reflect.Value          // receiver of methods for the service
	typ    reflect.Type           // type of the receiver
	method map[string]*methodType // registered methods
}

type request struct {
	ServiceMethod string   // format: "Service.Method"
	Seq           uint64   // sequence number chosen by client
	next          *request // for free list in Server
}

type response struct {
	ServiceMethod string    // echoes that of the Request
	Seq           uint64    // echoes that of the request
	Error         string    // error, if any.
	next          *response // for free list in Server
}

// server represents an RPC Server.
type server struct {
	name         string
	mu           sync.Mutex // protects the serviceMap
	serviceMap   map[string]*service
	reqLock      sync.Mutex // protects freeReq
	freeReq      *request
	respLock     sync.Mutex // protects freeResp
	freeResp     *response
	hdlrWrappers []HandlerWrapper
}

// Is this an exported - upper case - name?
func isExported(name string) bool {
	rune, _ := utf8.DecodeRuneInString(name)
	return unicode.IsUpper(rune)
}

// Is this type exported or a builtin?
func isExportedOrBuiltinType(t reflect.Type) bool {
	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	// PkgPath will be non-empty even for an exported type,
	// so we need to check the type name as well.
	return isExported(t.Name()) || t.PkgPath() == ""
}

// prepareMethod returns a methodType for the provided method or nil
// in case if the method was unsuitable.
func prepareMethod(method reflect.Method) *methodType {
	mtype := method.Type
	mname := method.Name
	var replyType, argType, contextType reflect.Type
	var stream bool

	// Method must be exported.
	if method.PkgPath != "" {
		return nil
	}

	switch mtype.NumIn() {
	case 3:
		// assuming streaming
		argType = mtype.In(2)
		contextType = mtype.In(1)
		stream = true
	case 4:
		// method that takes a context
		argType = mtype.In(2)
		replyType = mtype.In(3)
		contextType = mtype.In(1)
	default:
		log.Println("method", mname, "of", mtype, "has wrong number of ins:", mtype.NumIn())
		return nil
	}

	if stream {
		// check stream type
		streamType := reflect.TypeOf((*Streamer)(nil)).Elem()
		if !argType.Implements(streamType) {
			log.Println(mname, "argument does not implement Streamer interface:", argType)
			return nil
		}
	} else {
		// if not stream check the replyType

		// First arg need not be a pointer.
		if !isExportedOrBuiltinType(argType) {
			log.Println(mname, "argument type not exported:", argType)
			return nil
		}

		if replyType.Kind() != reflect.Ptr {
			log.Println("method", mname, "reply type not a pointer:", replyType)
			return nil
		}

		// Reply type must be exported.
		if !isExportedOrBuiltinType(replyType) {
			log.Println("method", mname, "reply type not exported:", replyType)
			return nil
		}
	}

	// Method needs one out.
	if mtype.NumOut() != 1 {
		log.Println("method", mname, "has wrong number of outs:", mtype.NumOut())
		return nil
	}
	// The return type of the method must be error.
	if returnType := mtype.Out(0); returnType != typeOfError {
		log.Println("method", mname, "returns", returnType.String(), "not error")
		return nil
	}
	return &methodType{method: method, ArgType: argType, ReplyType: replyType, ContextType: contextType, stream: stream}
}

func (server *server) register(rcvr interface{}) error {
	server.mu.Lock()
	defer server.mu.Unlock()
	if server.serviceMap == nil {
		server.serviceMap = make(map[string]*service)
	}
	s := new(service)
	s.typ = reflect.TypeOf(rcvr)
	s.rcvr = reflect.ValueOf(rcvr)
	sname := reflect.Indirect(s.rcvr).Type().Name()
	if sname == "" {
		log.Fatal("rpc: no service name for type", s.typ.String())
	}
	if !isExported(sname) {
		s := "rpc Register: type " + sname + " is not exported"
		log.Print(s)
		return errors.New(s)
	}
	if _, present := server.serviceMap[sname]; present {
		return errors.New("rpc: service already defined: " + sname)
	}
	s.name = sname
	s.method = make(map[string]*methodType)

	// Install the methods
	for m := 0; m < s.typ.NumMethod(); m++ {
		method := s.typ.Method(m)
		if mt := prepareMethod(method); mt != nil {
			s.method[method.Name] = mt
		}
	}

	if len(s.method) == 0 {
		s := "rpc Register: type " + sname + " has no exported methods of suitable type"
		log.Print(s)
		return errors.New(s)
	}
	server.serviceMap[s.name] = s
	return nil
}

func (server *server) sendResponse(sending *sync.Mutex, req *request, reply interface{}, codec serverCodec, errmsg string, last bool) (err error) {
	resp := server.getResponse()
	// Encode the response header
	resp.ServiceMethod = req.ServiceMethod
	if errmsg != "" {
		resp.Error = errmsg
		reply = invalidRequest
	}
	resp.Seq = req.Seq
	sending.Lock()
	err = codec.WriteResponse(resp, reply, last)
	if err != nil {
		log.Println("rpc: writing response:", err)
	}
	sending.Unlock()
	server.freeResponse(resp)
	return err
}

func (s *service) call(ctx context.Context, server *server, sending *sync.Mutex, mtype *methodType, req *request, argv, replyv reflect.Value, codec serverCodec, ct string) {
	mtype.Lock()
	mtype.numCalls++
	mtype.Unlock()
	function := mtype.method.Func
	var returnValues []reflect.Value

	r := &rpcRequest{
		service:     server.name,
		contentType: ct,
		method:      req.ServiceMethod,
	}

	if !mtype.stream {
		r.request = argv.Interface()

		fn := func(ctx context.Context, req Request, rsp interface{}) error {
			returnValues = function.Call([]reflect.Value{s.rcvr, mtype.prepareContext(ctx), reflect.ValueOf(req.Request()), reflect.ValueOf(rsp)})

			// The return value for the method is an error.
			if err := returnValues[0].Interface(); err != nil {
				return err.(error)
			}

			return nil
		}

		for i := len(server.hdlrWrappers); i > 0; i-- {
			fn = server.hdlrWrappers[i-1](fn)
		}

		errmsg := ""
		err := fn(ctx, r, replyv.Interface())
		if err != nil {
			errmsg = err.Error()
		}

		server.sendResponse(sending, req, replyv.Interface(), codec, errmsg, true)
		server.freeRequest(req)
		return
	}

	// declare a local error to see if we errored out already
	// keep track of the type, to make sure we return
	// the same one consistently
	var lastError error

	stream := &rpcStream{
		context: ctx,
		codec:   codec,
		request: r,
		seq:     req.Seq,
	}

	// Invoke the method, providing a new value for the reply.
	fn := func(ctx context.Context, req Request, stream interface{}) error {
		returnValues = function.Call([]reflect.Value{s.rcvr, mtype.prepareContext(ctx), reflect.ValueOf(stream)})
		if err := returnValues[0].Interface(); err != nil {
			// the function returned an error, we use that
			return err.(error)
		} else if lastError != nil {
			// we had an error inside sendReply, we use that
			return lastError
		} else {
			// no error, we send the special EOS error
			return lastStreamResponseError
		}
	}

	for i := len(server.hdlrWrappers); i > 0; i-- {
		fn = server.hdlrWrappers[i-1](fn)
	}

	// client.Stream request
	r.stream = true

	errmsg := ""
	if err := fn(ctx, r, stream); err != nil {
		errmsg = err.Error()
	}

	// this is the last packet, we don't do anything with
	// the error here (well sendStreamResponse will log it
	// already)
	server.sendResponse(sending, req, nil, codec, errmsg, true)
	server.freeRequest(req)
}

func (m *methodType) prepareContext(ctx context.Context) reflect.Value {
	if contextv := reflect.ValueOf(ctx); contextv.IsValid() {
		return contextv
	}
	return reflect.Zero(m.ContextType)
}

func (server *server) serveRequest(ctx context.Context, codec serverCodec, ct string) error {
	sending := new(sync.Mutex)
	service, mtype, req, argv, replyv, keepReading, err := server.readRequest(codec)
	if err != nil {
		if !keepReading {
			return err
		}
		// send a response if we actually managed to read a header.
		if req != nil {
			server.sendResponse(sending, req, invalidRequest, codec, err.Error(), true)
			server.freeRequest(req)
		}
		return err
	}
	service.call(ctx, server, sending, mtype, req, argv, replyv, codec, ct)
	return nil
}

func (server *server) getRequest() *request {
	server.reqLock.Lock()
	req := server.freeReq
	if req == nil {
		req = new(request)
	} else {
		server.freeReq = req.next
		*req = request{}
	}
	server.reqLock.Unlock()
	return req
}

func (server *server) freeRequest(req *request) {
	server.reqLock.Lock()
	req.next = server.freeReq
	server.freeReq = req
	server.reqLock.Unlock()
}

func (server *server) getResponse() *response {
	server.respLock.Lock()
	resp := server.freeResp
	if resp == nil {
		resp = new(response)
	} else {
		server.freeResp = resp.next
		*resp = response{}
	}
	server.respLock.Unlock()
	return resp
}

func (server *server) freeResponse(resp *response) {
	server.respLock.Lock()
	resp.next = server.freeResp
	server.freeResp = resp
	server.respLock.Unlock()
}

func (server *server) readRequest(codec serverCodec) (service *service, mtype *methodType, req *request, argv, replyv reflect.Value, keepReading bool, err error) {
	service, mtype, req, keepReading, err = server.readRequestHeader(codec)
	if err != nil {
		if !keepReading {
			return
		}
		// discard body
		codec.ReadRequestBody(nil)
		return
	}
	// is it a streaming request? then we don't read the body
	if mtype.stream {
		codec.ReadRequestBody(nil)
		return
	}

	// Decode the argument value.
	argIsValue := false // if true, need to indirect before calling.
	if mtype.ArgType.Kind() == reflect.Ptr {
		argv = reflect.New(mtype.ArgType.Elem())
	} else {
		argv = reflect.New(mtype.ArgType)
		argIsValue = true
	}
	// argv guaranteed to be a pointer now.
	if err = codec.ReadRequestBody(argv.Interface()); err != nil {
		return
	}
	if argIsValue {
		argv = argv.Elem()
	}

	if !mtype.stream {
		replyv = reflect.New(mtype.ReplyType.Elem())
	}
	return
}

func (server *server) readRequestHeader(codec serverCodec) (service *service, mtype *methodType, req *request, keepReading bool, err error) {
	// Grab the request header.
	req = server.getRequest()
	err = codec.ReadRequestHeader(req, true)
	if err != nil {
		req = nil
		if err == io.EOF || err == io.ErrUnexpectedEOF {
			return
		}
		err = errors.New("rpc: server cannot decode request: " + err.Error())
		return
	}

	// We read the header successfully. If we see an error now,
	// we can still recover and move on to the next request.
	keepReading = true

	serviceMethod := strings.Split(req.ServiceMethod, ".")
	if len(serviceMethod) != 2 {
		err = errors.New("rpc: service/method request ill-formed: " + req.ServiceMethod)
		return
	}
	// Look up the request.
	server.mu.Lock()
	service = server.serviceMap[serviceMethod[0]]
	server.mu.Unlock()
	if service == nil {
		err = errors.New("rpc: can't find service " + req.ServiceMethod)
		return
	}
	mtype = service.method[serviceMethod[1]]
	if mtype == nil {
		err = errors.New("rpc: can't find method " + req.ServiceMethod)
	}
	return
}

type serverCodec interface {
	ReadRequestHeader(*request, bool) error
	ReadRequestBody(interface{}) error
	WriteResponse(*response, interface{}, bool) error

	Close() error
}
<commit_msg>We don't need numcalls, save us a mutex lock<commit_after>package server

// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Meh, we need to get rid of this shit

import (
	"errors"
	"io"
	"log"
	"reflect"
	"strings"
	"sync"
	"unicode"
	"unicode/utf8"

	"golang.org/x/net/context"
)

var (
	lastStreamResponseError = errors.New("EOS")
	// A value sent as a placeholder for the server's response value when the server
	// receives an invalid request. It is never decoded by the client since the Response
	// contains an error when it is used.
	invalidRequest = struct{}{}

	// Precompute the reflect type for error. Can't use error directly
	// because Typeof takes an empty interface value. This is annoying.
	typeOfError = reflect.TypeOf((*error)(nil)).Elem()
)

type methodType struct {
	sync.Mutex  // protects counters
	method      reflect.Method
	ArgType     reflect.Type
	ReplyType   reflect.Type
	ContextType reflect.Type
	stream      bool
}

type service struct {
	name   string                 // name of service
	rcvr   reflect.Value          // receiver of methods for the service
	typ    reflect.Type           // type of the receiver
	method map[string]*methodType // registered methods
}

type request struct {
	ServiceMethod string   // format: "Service.Method"
	Seq           uint64   // sequence number chosen by client
	next          *request // for free list in Server
}

type response struct {
	ServiceMethod string    // echoes that of the Request
	Seq           uint64    // echoes that of the request
	Error         string    // error, if any.
	next          *response // for free list in Server
}

// server represents an RPC Server.
type server struct {
	name         string
	mu           sync.Mutex // protects the serviceMap
	serviceMap   map[string]*service
	reqLock      sync.Mutex // protects freeReq
	freeReq      *request
	respLock     sync.Mutex // protects freeResp
	freeResp     *response
	hdlrWrappers []HandlerWrapper
}

// Is this an exported - upper case - name?
func isExported(name string) bool {
	rune, _ := utf8.DecodeRuneInString(name)
	return unicode.IsUpper(rune)
}

// Is this type exported or a builtin?
func isExportedOrBuiltinType(t reflect.Type) bool {
	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	// PkgPath will be non-empty even for an exported type,
	// so we need to check the type name as well.
	return isExported(t.Name()) || t.PkgPath() == ""
}

// prepareMethod returns a methodType for the provided method or nil
// in case if the method was unsuitable.
func prepareMethod(method reflect.Method) *methodType {
	mtype := method.Type
	mname := method.Name
	var replyType, argType, contextType reflect.Type
	var stream bool

	// Method must be exported.
	if method.PkgPath != "" {
		return nil
	}

	switch mtype.NumIn() {
	case 3:
		// assuming streaming
		argType = mtype.In(2)
		contextType = mtype.In(1)
		stream = true
	case 4:
		// method that takes a context
		argType = mtype.In(2)
		replyType = mtype.In(3)
		contextType = mtype.In(1)
	default:
		log.Println("method", mname, "of", mtype, "has wrong number of ins:", mtype.NumIn())
		return nil
	}

	if stream {
		// check stream type
		streamType := reflect.TypeOf((*Streamer)(nil)).Elem()
		if !argType.Implements(streamType) {
			log.Println(mname, "argument does not implement Streamer interface:", argType)
			return nil
		}
	} else {
		// if not stream check the replyType

		// First arg need not be a pointer.
		if !isExportedOrBuiltinType(argType) {
			log.Println(mname, "argument type not exported:", argType)
			return nil
		}

		if replyType.Kind() != reflect.Ptr {
			log.Println("method", mname, "reply type not a pointer:", replyType)
			return nil
		}

		// Reply type must be exported.
		if !isExportedOrBuiltinType(replyType) {
			log.Println("method", mname, "reply type not exported:", replyType)
			return nil
		}
	}

	// Method needs one out.
	if mtype.NumOut() != 1 {
		log.Println("method", mname, "has wrong number of outs:", mtype.NumOut())
		return nil
	}
	// The return type of the method must be error.
	if returnType := mtype.Out(0); returnType != typeOfError {
		log.Println("method", mname, "returns", returnType.String(), "not error")
		return nil
	}
	return &methodType{method: method, ArgType: argType, ReplyType: replyType, ContextType: contextType, stream: stream}
}

func (server *server) register(rcvr interface{}) error {
	server.mu.Lock()
	defer server.mu.Unlock()
	if server.serviceMap == nil {
		server.serviceMap = make(map[string]*service)
	}
	s := new(service)
	s.typ = reflect.TypeOf(rcvr)
	s.rcvr = reflect.ValueOf(rcvr)
	sname := reflect.Indirect(s.rcvr).Type().Name()
	if sname == "" {
		log.Fatal("rpc: no service name for type", s.typ.String())
	}
	if !isExported(sname) {
		s := "rpc Register: type " + sname + " is not exported"
		log.Print(s)
		return errors.New(s)
	}
	if _, present := server.serviceMap[sname]; present {
		return errors.New("rpc: service already defined: " + sname)
	}
	s.name = sname
	s.method = make(map[string]*methodType)

	// Install the methods
	for m := 0; m < s.typ.NumMethod(); m++ {
		method := s.typ.Method(m)
		if mt := prepareMethod(method); mt != nil {
			s.method[method.Name] = mt
		}
	}

	if len(s.method) == 0 {
		s := "rpc Register: type " + sname + " has no exported methods of suitable type"
		log.Print(s)
		return errors.New(s)
	}
	server.serviceMap[s.name] = s
	return nil
}

func (server *server) sendResponse(sending *sync.Mutex, req *request, reply interface{}, codec serverCodec, errmsg string, last bool) (err error) {
	resp := server.getResponse()
	// Encode the response header
	resp.ServiceMethod = req.ServiceMethod
	if errmsg != "" {
		resp.Error = errmsg
		reply = invalidRequest
	}
	resp.Seq = req.Seq
	sending.Lock()
	err = codec.WriteResponse(resp, reply, last)
	if err != nil {
		log.Println("rpc: writing response:", err)
	}
	sending.Unlock()
	server.freeResponse(resp)
	return err
}

func (s *service) call(ctx context.Context, server *server, sending *sync.Mutex, mtype *methodType, req *request, argv, replyv reflect.Value, codec serverCodec, ct string) {
	function := mtype.method.Func
	var returnValues []reflect.Value

	r := &rpcRequest{
		service:     server.name,
		contentType: ct,
		method:      req.ServiceMethod,
	}

	if !mtype.stream {
		r.request = argv.Interface()

		fn := func(ctx context.Context, req Request, rsp interface{}) error {
			returnValues = function.Call([]reflect.Value{s.rcvr, mtype.prepareContext(ctx), reflect.ValueOf(req.Request()), reflect.ValueOf(rsp)})

			// The return value for the method is an error.
			if err := returnValues[0].Interface(); err != nil {
				return err.(error)
			}

			return nil
		}

		for i := len(server.hdlrWrappers); i > 0; i-- {
			fn = server.hdlrWrappers[i-1](fn)
		}

		errmsg := ""
		err := fn(ctx, r, replyv.Interface())
		if err != nil {
			errmsg = err.Error()
		}

		server.sendResponse(sending, req, replyv.Interface(), codec, errmsg, true)
		server.freeRequest(req)
		return
	}

	// declare a local error to see if we errored out already
	// keep track of the type, to make sure we return
	// the same one consistently
	var lastError error

	stream := &rpcStream{
		context: ctx,
		codec:   codec,
		request: r,
		seq:     req.Seq,
	}

	// Invoke the method, providing a new value for the reply.
	fn := func(ctx context.Context, req Request, stream interface{}) error {
		returnValues = function.Call([]reflect.Value{s.rcvr, mtype.prepareContext(ctx), reflect.ValueOf(stream)})
		if err := returnValues[0].Interface(); err != nil {
			// the function returned an error, we use that
			return err.(error)
		} else if lastError != nil {
			// we had an error inside sendReply, we use that
			return lastError
		} else {
			// no error, we send the special EOS error
			return lastStreamResponseError
		}
	}

	for i := len(server.hdlrWrappers); i > 0; i-- {
		fn = server.hdlrWrappers[i-1](fn)
	}

	// client.Stream request
	r.stream = true

	errmsg := ""
	if err := fn(ctx, r, stream); err != nil {
		errmsg = err.Error()
	}

	// this is the last packet, we don't do anything with
	// the error here (well sendStreamResponse will log it
	// already)
	server.sendResponse(sending, req, nil, codec, errmsg, true)
	server.freeRequest(req)
}

func (m *methodType) prepareContext(ctx context.Context) reflect.Value {
	if contextv := reflect.ValueOf(ctx); contextv.IsValid() {
		return contextv
	}
	return reflect.Zero(m.ContextType)
}

func (server *server) serveRequest(ctx context.Context, codec serverCodec, ct string) error {
	sending := new(sync.Mutex)
	service, mtype, req, argv, replyv, keepReading, err := server.readRequest(codec)
	if err != nil {
		if !keepReading {
			return err
		}
		// send a response if we actually managed to read a header.
		if req != nil {
			server.sendResponse(sending, req, invalidRequest, codec, err.Error(), true)
			server.freeRequest(req)
		}
		return err
	}
	service.call(ctx, server, sending, mtype, req, argv, replyv, codec, ct)
	return nil
}

func (server *server) getRequest() *request {
	server.reqLock.Lock()
	req := server.freeReq
	if req == nil {
		req = new(request)
	} else {
		server.freeReq = req.next
		*req = request{}
	}
	server.reqLock.Unlock()
	return req
}

func (server *server) freeRequest(req *request) {
	server.reqLock.Lock()
	req.next = server.freeReq
	server.freeReq = req
	server.reqLock.Unlock()
}

func (server *server) getResponse() *response {
	server.respLock.Lock()
	resp := server.freeResp
	if resp == nil {
		resp = new(response)
	} else {
		server.freeResp = resp.next
		*resp = response{}
	}
	server.respLock.Unlock()
	return resp
}

func (server *server) freeResponse(resp *response) {
	server.respLock.Lock()
	resp.next = server.freeResp
	server.freeResp = resp
	server.respLock.Unlock()
}

func (server *server) readRequest(codec serverCodec) (service *service, mtype *methodType, req *request, argv, replyv reflect.Value, keepReading bool, err error) {
	service, mtype, req, keepReading, err = server.readRequestHeader(codec)
	if err != nil {
		if !keepReading {
			return
		}
		// discard body
		codec.ReadRequestBody(nil)
		return
	}
	// is it a streaming request? then we don't read the body
	if mtype.stream {
		codec.ReadRequestBody(nil)
		return
	}

	// Decode the argument value.
	argIsValue := false // if true, need to indirect before calling.
	if mtype.ArgType.Kind() == reflect.Ptr {
		argv = reflect.New(mtype.ArgType.Elem())
	} else {
		argv = reflect.New(mtype.ArgType)
		argIsValue = true
	}
	// argv guaranteed to be a pointer now.
	if err = codec.ReadRequestBody(argv.Interface()); err != nil {
		return
	}
	if argIsValue {
		argv = argv.Elem()
	}

	if !mtype.stream {
		replyv = reflect.New(mtype.ReplyType.Elem())
	}
	return
}

func (server *server) readRequestHeader(codec serverCodec) (service *service, mtype *methodType, req *request, keepReading bool, err error) {
	// Grab the request header.
	req = server.getRequest()
	err = codec.ReadRequestHeader(req, true)
	if err != nil {
		req = nil
		if err == io.EOF || err == io.ErrUnexpectedEOF {
			return
		}
		err = errors.New("rpc: server cannot decode request: " + err.Error())
		return
	}

	// We read the header successfully. If we see an error now,
	// we can still recover and move on to the next request.
	keepReading = true

	serviceMethod := strings.Split(req.ServiceMethod, ".")
	if len(serviceMethod) != 2 {
		err = errors.New("rpc: service/method request ill-formed: " + req.ServiceMethod)
		return
	}
	// Look up the request.
	server.mu.Lock()
	service = server.serviceMap[serviceMethod[0]]
	server.mu.Unlock()
	if service == nil {
		err = errors.New("rpc: can't find service " + req.ServiceMethod)
		return
	}
	mtype = service.method[serviceMethod[1]]
	if mtype == nil {
		err = errors.New("rpc: can't find method " + req.ServiceMethod)
	}
	return
}

type serverCodec interface {
	ReadRequestHeader(*request, bool) error
	ReadRequestBody(interface{}) error
	WriteResponse(*response, interface{}, bool) error

	Close() error
}
<|endoftext|>


<commit_before>package main

import (
	"common"
	"common/crypto"
	"common/erasure"
	"network"
	"testing"
)

type ServerHandler struct {
	seg common.Segment
}

func (sh *ServerHandler) UploadSegment(seg common.Segment, arb *struct{}) error {
	sh.seg = seg
	return nil
}

func (sh *ServerHandler) DownloadSegment(hash crypto.Hash, seg *common.Segment) error {
	*seg = sh.seg
	return nil
}

// TestRPCUploadSector tests the NewRPCServer and uploadFile functions.
// NewRPCServer must properly initialize a RPC server.
// uploadSector must succesfully distribute a Sector among a quorum.
// The uploaded Sector must be successfully reconstructed.
func TestRPCuploadSector(t *testing.T) {
	// create RPCServer
	rpcs, err := network.NewRPCServer(9985)
	if err != nil {
		t.Fatal("Failed to initialize RPCServer:", err)
	}
	defer rpcs.Close()

	// create quorum
	var q [common.QuorumSize]common.Address
	var shs [common.QuorumSize]ServerHandler
	for i := 0; i < common.QuorumSize; i++ {
		q[i] = common.Address{0, "localhost", 9000 + i}
		qrpc, err := network.NewRPCServer(9000 + i)
		defer qrpc.Close()
		if err != nil {
			t.Fatal("Failed to initialize RPCServer:", err)
		}
		q[i].ID = qrpc.RegisterHandler(&shs[i])
	}

	// create sector
	secData, err := crypto.RandomByteSlice(70000)
	if err != nil {
		t.Fatal("Could not generate test data:", err)
	}

	sec, err := common.NewSector(secData)
	if err != nil {
		t.Fatal("Failed to create sector:", err)
	}

	// upload sector to quorum
	k := common.QuorumSize / 2
	ring, err := uploadSector(rpcs, sec, k, q)
	if err != nil {
		t.Fatal("Failed to upload file:", err)
	}

	// rebuild file from first k segments
	ring.Segs = []common.Segment{}
	for i := 0; i < k; i++ {
		ring.AddSegment(&shs[i].seg)
	}

	sec, err = erasure.RebuildSector(ring)
	if err != nil {
		t.Fatal("Failed to rebuild file:", err)
	}

	// check hash
	rebuiltHash, err := crypto.CalculateHash(sec.Data)
	if err != nil {
		t.Fatal("Failed to calculate hash:", err)
	}

	if sec.Hash != rebuiltHash {
		t.Fatal("Failed to recover file: hashes do not match")
	}
}

// TestRPCdownloadSector tests the NewRPCServer and downloadSector functions.
// NewRPCServer must properly initialize a RPC server.
// downloadSector must successfully retrieve a Sector from a quorum.
// The downloaded Sector must match the original Sector.
func TestRPCdownloadSector(t *testing.T) {
	// create sector
	secData, err := crypto.RandomByteSlice(70000)
	if err != nil {
		t.Fatal("Could not generate test data:", err)
	}

	sec, err := common.NewSector(secData)
	if err != nil {
		t.Fatal("Failed to create sector:", err)
	}

	// encode sector
	k := common.QuorumSize / 2
	ring, err := erasure.EncodeRing(sec, k)
	if err != nil {
		t.Fatal("Failed to encode sector data:", err)
	}

	// create RPCServer
	rpcs, err := network.NewRPCServer(9985)
	if err != nil {
		t.Fatal("Failed to initialize RPCServer:", err)
	}
	defer rpcs.Close()

	// create quorum
	var q [common.QuorumSize]common.Address
	for i := 0; i < common.QuorumSize; i++ {
		q[i] = common.Address{0, "localhost", 9000 + i}
		qrpc, err := network.NewRPCServer(9000 + i)
		if err != nil {
			t.Fatal("Failed to initialize RPCServer:", err)
		}
		sh := new(ServerHandler)
		sh.seg = ring.Segs[i]
		q[i].ID = qrpc.RegisterHandler(sh)
	}

	// download file from quorum
	ring.Segs = []common.Segment{}
	sec, err = downloadSector(rpcs, sec.Hash, ring, q)
	if err != nil {
		t.Fatal("Failed to download file:", err)
	}

	// check hash
	rebuiltHash, err := crypto.CalculateHash(sec.Data)
	if err != nil {
		t.Fatal("Failed to calculate hash:", err)
	}

	if sec.Hash != rebuiltHash {
		t.Fatal("Failed to recover file: hashes do not match")
	}
}
<commit_msg>update client test<commit_after>package main

import (
	"common"
	"common/crypto"
	"common/erasure"
	"network"
	"testing"
)

type Server struct {
	seg common.Segment
}

func (s *Server) UploadSegment(seg common.Segment, arb *struct{}) error {
	s.seg = seg
	return nil
}

func (s *Server) DownloadSegment(hash crypto.Hash, seg *common.Segment) error {
	*seg = s.seg
	return nil
}

// TestRPCUploadSector tests the NewRPCServer and uploadFile functions.
// NewRPCServer must properly initialize a RPC server.
// uploadSector must succesfully distribute a Sector among a quorum.
// The uploaded Sector must be successfully reconstructed.
func TestRPCuploadSector(t *testing.T) {
	SectorDB = make(map[crypto.Hash]*common.Ring)
	// create RPCServer
	var err error
	mr, err = network.NewRPCServer(9985)
	if err != nil {
		t.Fatal("Failed to initialize RPCServer:", err)
	}
	defer mr.Close()

	// create quorum
	var q [common.QuorumSize]common.Address
	var shs [common.QuorumSize]Server
	for i := 0; i < common.QuorumSize; i++ {
		q[i] = common.Address{0, "localhost", 9000 + i}
		qrpc, err := network.NewRPCServer(9000 + i)
		defer qrpc.Close()
		if err != nil {
			t.Fatal("Failed to initialize RPCServer:", err)
		}
		q[i].ID = qrpc.RegisterHandler(&shs[i])
	}

	// create sector
	secData, err := crypto.RandomByteSlice(70000)
	if err != nil {
		t.Fatal("Could not generate test data:", err)
	}

	sec, err := common.NewSector(secData)
	if err != nil {
		t.Fatal("Failed to create sector:", err)
	}

	// upload sector to quorum
	k := common.QuorumSize / 2
	err = uploadSector(sec, k, q)
	if err != nil {
		t.Fatal("Failed to upload file:", err)
	}

	// rebuild file from first k segments
	newRing := SectorDB[sec.Hash]
	var newSegs []common.Segment
	for i := 0; i < k; i++ {
		newSegs = append(newSegs, shs[i].seg)
	}

	sec, err = erasure.RebuildSector(newRing, newSegs)
	if err != nil {
		t.Fatal("Failed to rebuild file:", err)
	}

	// check hash
	rebuiltHash, err := crypto.CalculateHash(sec.Data)
	if err != nil {
		t.Fatal("Failed to calculate hash:", err)
	}

	if sec.Hash != rebuiltHash {
		t.Fatal("Failed to recover file: hashes do not match")
	}
}

// TestRPCdownloadSector tests the NewRPCServer and downloadSector functions.
// NewRPCServer must properly initialize a RPC server.
// downloadSector must successfully retrieve a Sector from a quorum.
// The downloaded Sector must match the original Sector.
func TestRPCdownloadSector(t *testing.T) {
	SectorDB = make(map[crypto.Hash]*common.Ring)
	// create sector
	secData, err := crypto.RandomByteSlice(70000)
	if err != nil {
		t.Fatal("Could not generate test data:", err)
	}

	sec, err := common.NewSector(secData)
	if err != nil {
		t.Fatal("Failed to create sector:", err)
	}

	// encode sector
	k := common.QuorumSize / 2
	ring, segs, err := erasure.EncodeRing(sec, k)
	if err != nil {
		t.Fatal("Failed to encode sector data:", err)
	}

	// create RPCServer
	mr, err = network.NewRPCServer(9985)
	if err != nil {
		t.Fatal("Failed to initialize RPCServer:", err)
	}
	defer mr.Close()

	// create quorum
	var q [common.QuorumSize]common.Address
	for i := 0; i < common.QuorumSize; i++ {
		q[i] = common.Address{0, "localhost", 9000 + i}
		qrpc, err := network.NewRPCServer(9000 + i)
		if err != nil {
			t.Fatal("Failed to initialize RPCServer:", err)
		}
		sh := new(Server)
		sh.seg = segs[i]
		q[i].ID = qrpc.RegisterHandler(sh)
	}

	// manually add entry to SectorDB
	ring.Hosts = q
	SectorDB[sec.Hash] = ring

	// download file from quorum
	sec, err = downloadSector(sec.Hash)
	if err != nil {
		t.Fatal("Failed to download file:", err)
	}

	// check hash
	rebuiltHash, err := crypto.CalculateHash(sec.Data)
	if err != nil {
		t.Fatal("Failed to calculate hash:", err)
	}

	if sec.Hash != rebuiltHash {
		t.Fatal("Failed to recover file: hashes do not match")
	}
}
<|endoftext|>