<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/dfordsoft/golib/ebook"
	"github.com/dfordsoft/golib/ic"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"
	"time"
)

func init() {
	registerNovelSiteHandler(&NovelSiteHandler{
		Match:    isUUKanshu,
		Download: dlUUKanshu,
	})
}

func isUUKanshu(u string) bool {
	r, _ := regexp.Compile(`http://www\.uukanshu\.net/b/[0-9]+/`)
	if r.MatchString(u) {
		return true
	}
	return false
}

func dlUUKanshu(u string) {
	client := &http.Client{
		Timeout: 60 * time.Second,
	}
	retry := 0
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		log.Println("uukanshu - Could not parse novel request:", err)
		return
	}

	req.Header.Set("Referer", "http://www.uukanshu.net/")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:45.0) Gecko/20100101 Firefox/45.0")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("accept-language", `en-US,en;q=0.8`)
	req.Header.Set("Upgrade-Insecure-Requests", "1")
doRequest:
	resp, err := client.Do(req)
	if err != nil {
		log.Println("uukanshu - Could not send novel request:", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("uukanshu - novel request not 200")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("uukanshu - Reading response body failed", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	mobi := &ebook.Mobi{}
	mobi.Begin()

	var title string
	var lines []string
	//
第一千零八十五章 调教完毕……

	r, _ := regexp.Compile(`<li><a\shref="/b/[0-9]+/([0-9]+\.html)"\stitle="[^"]+"\starget="_blank">([^<]+)</a></li>$`)
	scanner := bufio.NewScanner(bytes.NewReader(b))
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		// convert from gbk to UTF-8
		l := ic.ConvertString("gbk", "utf-8", line)
		if title == "" {
			re, _ := regexp.Compile(`<h1\sid="timu">([^<]+)</h1>$`)
			ss := re.FindAllStringSubmatch(l, -1)
			if len(ss) > 0 && len(ss[0]) > 0 {
				s := ss[0]
				title = s[1]
				mobi.SetTitle(title)
				continue
			}
		}
		if r.MatchString(l) {
			lines = append([]string{l}, lines...)
		}
	}
	lines = lines[:len(lines)-1]
	for _, l := range lines {
		ss := r.FindAllStringSubmatch(l, -1)
		s := ss[0]
		finalURL := fmt.Sprintf("%s%s", u, s[1])
		c := dlUUKanshuPage(finalURL)
		mobi.AppendContent(s[2], finalURL, string(c))
		fmt.Println(s[2], finalURL, len(c), "bytes")
	}
	mobi.End()
}

func dlUUKanshuPage(u string) (c []byte) {
	client := &http.Client{
		Timeout: 60 * time.Second,
	}
	retry := 0
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		log.Println("uukanshu - Could not parse novel page request:", err)
		return
	}

	req.Header.Set("Referer", "http://www.uukanshu.net/")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:45.0) Gecko/20100101 Firefox/45.0")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("accept-language", `en-US,en;q=0.8`)
	req.Header.Set("Upgrade-Insecure-Requests", "1")
doRequest:
	resp, err := client.Do(req)
	if err != nil {
		log.Println("uukanshu - Could not send novel page request:", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("uukanshu - novel page request not 200")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	c, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("uukanshu - novel page content reading failed")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}
	c = ic.Convert("gbk", "utf-8", c)
	c = bytes.Replace(c, []byte("\r\n"), []byte(""), -1)
	c = bytes.Replace(c, []byte("\r"), []byte(""), -1)
	c = bytes.Replace(c, []byte("\n"), []byte(""), -1)
	idx := bytes.Index(c, []byte(""))
	if idx > 1 {
		c = c[idx:]
	}
	idx = bytes.Index(c, []byte(`</div>`))
	if idx > 1 {
		c = c[idx+6:]
	}
	startStr := []byte("
")
	idx = bytes.Index(c, startStr)
	if idx > 1 {
		idxEnd := bytes.Index(c[idx:], []byte("
"))
		if idxEnd > 1 {
			b := c[idx:]
			c = append(c[:idx], b[idxEnd+6:]...)
		}
	}
	idx = bytes.Index(c, []byte("
"))
	if idx > 1 {
		c = c[:idx]
	}
	c = bytes.Replace(c, []byte("

    "), []byte("

"), -1)
	c = bytes.Replace(c, []byte("    "), []byte(""), -1)
	return
}
<commit_msg>(*)fixed title of uukanshu book<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/dfordsoft/golib/ebook"
	"github.com/dfordsoft/golib/ic"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"
	"strings"
	"time"
)

func init() {
	registerNovelSiteHandler(&NovelSiteHandler{
		Match:    isUUKanshu,
		Download: dlUUKanshu,
	})
}

func isUUKanshu(u string) bool {
	r, _ := regexp.Compile(`http://www\.uukanshu\.net/b/[0-9]+/`)
	if r.MatchString(u) {
		return true
	}
	return false
}

func dlUUKanshu(u string) {
	client := &http.Client{
		Timeout: 60 * time.Second,
	}
	retry := 0
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		log.Println("uukanshu - Could not parse novel request:", err)
		return
	}

	req.Header.Set("Referer", "http://www.uukanshu.net/")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:45.0) Gecko/20100101 Firefox/45.0")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("accept-language", `en-US,en;q=0.8`)
	req.Header.Set("Upgrade-Insecure-Requests", "1")
doRequest:
	resp, err := client.Do(req)
	if err != nil {
		log.Println("uukanshu - Could not send novel request:", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("uukanshu - novel request not 200")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("uukanshu - Reading response body failed", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	mobi := &ebook.Mobi{}
	mobi.Begin()

	var title string
	var lines []string
	//
第一千零八十五章 调教完毕……

	r, _ := regexp.Compile(`<li><a\shref="/b/[0-9]+/([0-9]+\.html)"\stitle="[^"]+"\starget="_blank">([^<]+)</a></li>$`)
	scanner := bufio.NewScanner(bytes.NewReader(b))
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		// convert from gbk to UTF-8
		l := ic.ConvertString("gbk", "utf-8", line)
		if title == "" {
			//
调教初唐最新章节

			re, _ := regexp.Compile(`<h1><a\shref="/b/[0-9]+/"\stitle="[^"]+">([^<]+)</a></h1>$`)
			ss := re.FindAllStringSubmatch(l, -1)
			if len(ss) > 0 && len(ss[0]) > 0 {
				s := ss[0]
				title = s[1]
				idx := strings.Index(title, `最新章节`)
				if idx > 0 {
					title = title[:idx]
				}
				mobi.SetTitle(title)
				continue
			}
		}
		if r.MatchString(l) {
			lines = append([]string{l}, lines...)
		}
	}
	lines = lines[:len(lines)-1]
	for _, l := range lines {
		ss := r.FindAllStringSubmatch(l, -1)
		s := ss[0]
		finalURL := fmt.Sprintf("%s%s", u, s[1])
		c := dlUUKanshuPage(finalURL)
		mobi.AppendContent(s[2], finalURL, string(c))
		fmt.Println(s[2], finalURL, len(c), "bytes")
	}
	mobi.End()
}

func dlUUKanshuPage(u string) (c []byte) {
	client := &http.Client{
		Timeout: 60 * time.Second,
	}
	retry := 0
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		log.Println("uukanshu - Could not parse novel page request:", err)
		return
	}

	req.Header.Set("Referer", "http://www.uukanshu.net/")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:45.0) Gecko/20100101 Firefox/45.0")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("accept-language", `en-US,en;q=0.8`)
	req.Header.Set("Upgrade-Insecure-Requests", "1")
doRequest:
	resp, err := client.Do(req)
	if err != nil {
		log.Println("uukanshu - Could not send novel page request:", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("uukanshu - novel page request not 200")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	c, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("uukanshu - novel page content reading failed")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}
	c = ic.Convert("gbk", "utf-8", c)
	c = bytes.Replace(c, []byte("\r\n"), []byte(""), -1)
	c = bytes.Replace(c, []byte("\r"), []byte(""), -1)
	c = bytes.Replace(c, []byte("\n"), []byte(""), -1)
	idx := bytes.Index(c, []byte(""))
	if idx > 1 {
		c = c[idx:]
	}
	idx = bytes.Index(c, []byte(`</div>`))
	if idx > 1 {
		c = c[idx+6:]
	}
	startStr := []byte("
")
	idx = bytes.Index(c, startStr)
	if idx > 1 {
		idxEnd := bytes.Index(c[idx:], []byte("
"))
		if idxEnd > 1 {
			b := c[idx:]
			c = append(c[:idx], b[idxEnd+6:]...)
		}
	}
	idx = bytes.Index(c, []byte("
"))
	if idx > 1 {
		c = c[:idx]
	}
	c = bytes.Replace(c, []byte("

    "), []byte("

"), -1)
	c = bytes.Replace(c, []byte("    "), []byte(""), -1)
	c = bytes.Replace(c, []byte("
　　"), []byte("
"), -1)
	return
}
<|endoftext|>


<commit_before>package generator

import (
	"math"
)

const (
	// Frequency
	hz = .2

	// Amplitude (peak-to-peak)
	amp = 128

	aggSz = 2
)

func GenSineData(tuples []*SensorTuple) error {
	peakAmp := float64(amp / 2)

	for i := 0; i < len(tuples); i++ {
		tuples[i].Data = peakAmp * (math.Sin(hz*float64(i)) + 1.0)
		tuples[i].Data = toFixed(tuples[i].Data, 2)
	}

	return nil
}

func CalcAvg(tuples []*SensorTuple) error {
	sum := 0.0

	for i := 1; i <= len(tuples); i++ {
		sum += tuples[i-1].Data
		if i%aggSz == 0 {
			tuples[i-1].Aggregate = sum / float64(aggSz)
			tuples[i-1].Aggregate = toFixed(tuples[i-1].Aggregate, 2)
			sum = 0.0
		}
	}

	return nil
}

func toFixed(num float64, precision int) float64 {
	output := math.Pow(10, float64(precision))
	return float64(round(num*output)) / output
}

func round(num float64) int {
	return int(num + math.Copysign(0.5, num))
}
<commit_msg>Changed name of decimal rounding function<commit_after>package generator

import (
	"math"
)

const (
	// Frequency
	hz = .2

	// Amplitude (peak-to-peak)
	amp = 128

	aggSz = 2
)

func GenSineData(tuples []*SensorTuple) error {
	peakAmp := float64(amp / 2)

	for i := 0; i < len(tuples); i++ {
		tuples[i].Data = peakAmp * (math.Sin(hz*float64(i)) + 1.0)
		tuples[i].Data = roundDecimal(tuples[i].Data, 2)
	}

	return nil
}

func CalcAvg(tuples []*SensorTuple) error {
	sum := 0.0

	for i := 1; i <= len(tuples); i++ {
		sum += tuples[i-1].Data
		if i%aggSz == 0 {
			tuples[i-1].Aggregate = sum / float64(aggSz)
			tuples[i-1].Aggregate = roundDecimal(tuples[i-1].Aggregate, 2)
			sum = 0.0
		}
	}

	return nil
}

func roundDecimal(num float64, precision int) float64 {
	output := math.Pow(10, float64(precision))
	return float64(round(num*output)) / output
}

func round(num float64) int {
	return int(num + math.Copysign(0.5, num))
}
<|endoftext|>


<commit_before>// +build darwin dragonfly freebsd linux nacl netbsd openbsd solaris

package main

import "fmt"

// Users that are known on a Unix system to try using
var knownUsers = []string{"", "root"}
<commit_msg>Remove fmt import from test<commit_after>// +build darwin dragonfly freebsd linux nacl netbsd openbsd solaris

package main

// Users that are known on a Unix system to try using
var knownUsers = []string{"", "root"}
<|endoftext|>


<commit_before>package config // import "a4.io/blobstash/pkg/config"

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/inconshreveable/log15"
	"gopkg.in/yaml.v2"

	"a4.io/blobstash/pkg/config/pathutil"
)

var (
	DefaultListen  = ":8051"
	LetsEncryptDir = "letsencrypt"
)

// AppConfig holds an app configuration items
type AppConfig struct {
	Name       string `yaml:"name"`
	Path       string `yaml:"path"` // App path, optional?
	Entrypoint string `yaml:"entrypoint"`
	Domain     string `yaml:"domain"`
	Username   string `yaml:"username"`
	Password   string `yaml:"password"`
	Proxy      string `yaml:"proxy"`
	Remote     string `yaml:"remote"`
	Scheduled  string `yaml:"scheduled"`

	Config map[string]interface{} `yaml:"config"`
}

type S3Repl struct {
	Bucket    string `yaml:"bucket"`
	Region    string `yaml:"region"`
	KeyFile   string `yaml:"key_file"`
	Endpoint  string `yaml:"endpoint"`
	AccessKey string `yaml:"access_key_id"`
	SecretKey string `yaml:"secret_access_key"`
}

type Replication struct {
	EnableOplog bool `yaml:"enable_oplog"`
}

type ReplicateFrom struct {
	URL    string `yaml:"url"`
	APIKey string `yaml:"api_key"`
}

func (s3 *S3Repl) Key() (*[32]byte, error) {
	if s3.KeyFile == "" {
		return nil, nil
	}
	var out [32]byte
	data, err := ioutil.ReadFile(s3.KeyFile)
	if err != nil {
		return nil, err
	}
	copy(out[:], data)
	return &out, nil
}

type GitServerConfig struct {
	Namespaces map[string]*GitNamespaceConf `yaml:"namespaces"`
}

type GitNamespaceConf struct {
	Username string `yaml:"username"`
	Password string `yaml:"password"`
}

type BasicAuth struct {
	ID       string   `yaml:"id"`
	Roles    []string `yaml:"roles"`
	Username string   `yaml:"username"`
	Password string   `yaml:"password"`
}

type Role struct {
	Name     string                 `yaml:"name"`
	Template string                 `yaml:"template"`
	Perms    []*Perm                `yaml:"permissions'`
	Args     map[string]interface{} `yaml:"args"`

	// Only set pragmatically for "managed role"
	Managed      bool     `yaml:"-"`
	ArgsRequired []string `yaml:"-"`
}

type Perm struct {
	Action   string `yaml:"action"`
	Resource string `yaml:"resource"`
}

// Config holds the configuration items
type Config struct {
	init     bool
	Listen   string `yaml:"listen"`
	LogLevel string `yaml:"log_level"`
	// TLS     bool     `yaml:"tls"`
	AutoTLS bool     `yaml:"tls_auto"`
	Domains []string `yaml:"tls_domains"`

	Roles []*Role `yaml:"roles"`
	Auth  []*BasicAuth

	ExpvarListen string `yaml:"expvar_server_listen"`

	ExtraApacheCombinedLogs string `yaml:"extra_apache_combined_logs"`

	GitServer *GitServerConfig `yaml:"git_server"`

	SharingKey string  `yaml:"sharing_key"`
	DataDir    string  `yaml:"data_dir"`
	S3Repl     *S3Repl `yaml:"s3_replication"`

	Apps          []*AppConfig    `yaml:"apps"`
	Docstore      *DocstoreConfig `yaml:"docstore"`
	Replication   *Replication    `yaml:"replication"`
	ReplicateFrom *ReplicateFrom  `yaml:"replicate_from"`

	// Items defined with the CLI flags
	ScanMode      bool `yaml:"-"`
	S3ScanMode    bool `yaml:"-"`
	S3RestoreMode bool `yaml:"-"`
}

func (c *Config) LogLvl() log15.Lvl {
	if c.LogLevel == "" {
		c.LogLevel = "info"
	}
	lvl, err := log15.LvlFromString(c.LogLevel)
	if err != nil {
		panic(err)
	}
	return lvl
}

type DocstoreConfig struct {
	StoredQueries []*StoredQuery               `yaml:"stored_queries"`
	Hooks         map[string]map[string]string `yaml:"hooks"`
}

type StoredQuery struct {
	Name string `yaml:"name"`
	Path string `yaml:"path"`
}

// New initialize a config object by loading the YAML path at the given path
func New(path string) (*Config, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	conf := &Config{}
	if err := yaml.Unmarshal([]byte(data), &conf); err != nil {
		return nil, err
	}
	return conf, nil
}

// VarDir returns the directory where the index will be stored
func (c *Config) ConfigDir() string {
	// TODO(tsileo): allow override?
	return pathutil.ConfigDir()
}

// VarDir returns the directory where the index will be stored
func (c *Config) VarDir() string {
	if c.DataDir != "" {
		return c.DataDir
	}
	return pathutil.VarDir()
}

// VarDir returns the directory where the index will be stored
func (c *Config) StashDir() string {
	return filepath.Join(c.VarDir(), "stash")
}

// Init initialize the config.
//
// It will try to create all the needed directory.
func (c *Config) Init() error {
	if c.init {
		return nil
	}
	if _, err := os.Stat(c.VarDir()); os.IsNotExist(err) {
		if err := os.MkdirAll(c.VarDir(), 0700); err != nil {
			return err
		}
	}
	if _, err := os.Stat(c.StashDir()); os.IsNotExist(err) {
		if err := os.MkdirAll(c.VarDir(), 0700); err != nil {
			return err
		}
	}
	if _, err := os.Stat(c.ConfigDir()); os.IsNotExist(err) {
		if err := os.MkdirAll(c.ConfigDir(), 0700); err != nil {
			return err
		}
	}
	if _, err := os.Stat(filepath.Join(c.ConfigDir(), LetsEncryptDir)); os.IsNotExist(err) {
		if err := os.MkdirAll(filepath.Join(c.ConfigDir(), LetsEncryptDir), 0700); err != nil {
			return err
		}
	}
	if c.SharingKey == "" {
		return fmt.Errorf("missing `sharing_key` config item")
	}
	if c.S3Repl != nil {
		// Set default region
		if c.S3Repl.Region == "" {
			c.S3Repl.Region = "us-east-1"
		}
	}
	c.init = true
	return nil
}

// Sync url config parsing
//u, err := url.Parse("http://:123@127.0.0.1:8053")
//	if err != nil {
//		log.Fatal(err)
//	}
//	u.User = nil
//	//apiKey, _ := u.User.Password()
//	fmt.Printf("%+v", u)
<commit_msg>config: remove old config items<commit_after>package config // import "a4.io/blobstash/pkg/config"

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/inconshreveable/log15"
	"gopkg.in/yaml.v2"

	"a4.io/blobstash/pkg/config/pathutil"
)

var (
	DefaultListen  = ":8051"
	LetsEncryptDir = "letsencrypt"
)

// AppConfig holds an app configuration items
type AppConfig struct {
	Name       string `yaml:"name"`
	Path       string `yaml:"path"` // App path, optional?
	Entrypoint string `yaml:"entrypoint"`
	Domain     string `yaml:"domain"`
	Username   string `yaml:"username"`
	Password   string `yaml:"password"`
	Proxy      string `yaml:"proxy"`
	Remote     string `yaml:"remote"`
	Scheduled  string `yaml:"scheduled"`

	Config map[string]interface{} `yaml:"config"`
}

type S3Repl struct {
	Bucket    string `yaml:"bucket"`
	Region    string `yaml:"region"`
	KeyFile   string `yaml:"key_file"`
	Endpoint  string `yaml:"endpoint"`
	AccessKey string `yaml:"access_key_id"`
	SecretKey string `yaml:"secret_access_key"`
}

type Replication struct {
	EnableOplog bool `yaml:"enable_oplog"`
}

type ReplicateFrom struct {
	URL    string `yaml:"url"`
	APIKey string `yaml:"api_key"`
}

func (s3 *S3Repl) Key() (*[32]byte, error) {
	if s3.KeyFile == "" {
		return nil, nil
	}
	var out [32]byte
	data, err := ioutil.ReadFile(s3.KeyFile)
	if err != nil {
		return nil, err
	}
	copy(out[:], data)
	return &out, nil
}

type BasicAuth struct {
	ID       string   `yaml:"id"`
	Roles    []string `yaml:"roles"`
	Username string   `yaml:"username"`
	Password string   `yaml:"password"`
}

type Role struct {
	Name     string                 `yaml:"name"`
	Template string                 `yaml:"template"`
	Perms    []*Perm                `yaml:"permissions'`
	Args     map[string]interface{} `yaml:"args"`

	// Only set pragmatically for "managed role"
	Managed      bool     `yaml:"-"`
	ArgsRequired []string `yaml:"-"`
}

type Perm struct {
	Action   string `yaml:"action"`
	Resource string `yaml:"resource"`
}

// Config holds the configuration items
type Config struct {
	init     bool
	Listen   string `yaml:"listen"`
	LogLevel string `yaml:"log_level"`
	// TLS     bool     `yaml:"tls"`
	AutoTLS bool     `yaml:"tls_auto"`
	Domains []string `yaml:"tls_domains"`

	Roles []*Role `yaml:"roles"`
	Auth  []*BasicAuth

	ExpvarListen string `yaml:"expvar_server_listen"`

	ExtraApacheCombinedLogs string `yaml:"extra_apache_combined_logs"`

	GitServer *GitServerConfig `yaml:"git_server"`

	SharingKey string  `yaml:"sharing_key"`
	DataDir    string  `yaml:"data_dir"`
	S3Repl     *S3Repl `yaml:"s3_replication"`

	Apps          []*AppConfig    `yaml:"apps"`
	Docstore      *DocstoreConfig `yaml:"docstore"`
	Replication   *Replication    `yaml:"replication"`
	ReplicateFrom *ReplicateFrom  `yaml:"replicate_from"`

	// Items defined with the CLI flags
	ScanMode      bool `yaml:"-"`
	S3ScanMode    bool `yaml:"-"`
	S3RestoreMode bool `yaml:"-"`
}

func (c *Config) LogLvl() log15.Lvl {
	if c.LogLevel == "" {
		c.LogLevel = "info"
	}
	lvl, err := log15.LvlFromString(c.LogLevel)
	if err != nil {
		panic(err)
	}
	return lvl
}

type DocstoreConfig struct {
	StoredQueries []*StoredQuery               `yaml:"stored_queries"`
	Hooks         map[string]map[string]string `yaml:"hooks"`
}

type StoredQuery struct {
	Name string `yaml:"name"`
	Path string `yaml:"path"`
}

// New initialize a config object by loading the YAML path at the given path
func New(path string) (*Config, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	conf := &Config{}
	if err := yaml.Unmarshal([]byte(data), &conf); err != nil {
		return nil, err
	}
	return conf, nil
}

// VarDir returns the directory where the index will be stored
func (c *Config) ConfigDir() string {
	// TODO(tsileo): allow override?
	return pathutil.ConfigDir()
}

// VarDir returns the directory where the index will be stored
func (c *Config) VarDir() string {
	if c.DataDir != "" {
		return c.DataDir
	}
	return pathutil.VarDir()
}

// VarDir returns the directory where the index will be stored
func (c *Config) StashDir() string {
	return filepath.Join(c.VarDir(), "stash")
}

// Init initialize the config.
//
// It will try to create all the needed directory.
func (c *Config) Init() error {
	if c.init {
		return nil
	}
	if _, err := os.Stat(c.VarDir()); os.IsNotExist(err) {
		if err := os.MkdirAll(c.VarDir(), 0700); err != nil {
			return err
		}
	}
	if _, err := os.Stat(c.StashDir()); os.IsNotExist(err) {
		if err := os.MkdirAll(c.VarDir(), 0700); err != nil {
			return err
		}
	}
	if _, err := os.Stat(c.ConfigDir()); os.IsNotExist(err) {
		if err := os.MkdirAll(c.ConfigDir(), 0700); err != nil {
			return err
		}
	}
	if _, err := os.Stat(filepath.Join(c.ConfigDir(), LetsEncryptDir)); os.IsNotExist(err) {
		if err := os.MkdirAll(filepath.Join(c.ConfigDir(), LetsEncryptDir), 0700); err != nil {
			return err
		}
	}
	if c.SharingKey == "" {
		return fmt.Errorf("missing `sharing_key` config item")
	}
	if c.S3Repl != nil {
		// Set default region
		if c.S3Repl.Region == "" {
			c.S3Repl.Region = "us-east-1"
		}
	}
	c.init = true
	return nil
}

// Sync url config parsing
//u, err := url.Parse("http://:123@127.0.0.1:8053")
//	if err != nil {
//		log.Fatal(err)
//	}
//	u.User = nil
//	//apiKey, _ := u.User.Password()
//	fmt.Printf("%+v", u)
<|endoftext|>


<commit_before>package origin

import (
	"io"

	"github.com/tsavola/gate/service"
)

const (
	Name    = "origin"
	Version = 0

	packetHeaderSize = 8

	maxPacketSize = 0x10000 // TODO: move this elsewhere
)

type Factory struct {
	R io.Reader
	W io.Writer
}

func (f *Factory) Register(r *service.Registry) {
	service.Register(r, Name, Version, f)
}

func (f *Factory) New() service.Instance {
	return &origin{r: f.R, w: f.W}
}

var Default = new(Factory)

func Register(r *service.Registry) {
	Default.Register(r)
}

func CloneRegistryWith(r *service.Registry, origIn io.Reader, origOut io.Writer) *service.Registry {
	clone := service.Clone(r)
	(&Factory{R: origIn, W: origOut}).Register(clone)
	return clone
}

type origin struct {
	r io.Reader
	w io.Writer

	reading chan struct{}
}

func (o *origin) Handle(buf []byte, replies chan<- []byte) {
	if o.r != nil && o.reading == nil {
		o.reading = make(chan struct{})
		go o.readLoop(buf[6:8], replies)
	}

	if o.w != nil {
		if _, err := o.w.Write(buf[packetHeaderSize:]); err != nil {
			// assume that the error is EOF, broken pipe or such
			o.w = nil
		}
	}
}

func (o *origin) Shutdown() {
	if o.reading != nil {
		close(o.reading)
	}
}

func (o *origin) readLoop(code []byte, replies chan<- []byte) {
	for {
		buf := make([]byte, maxPacketSize) // TODO: smaller buffer?
		copy(buf[6:8], code)

		n, err := o.r.Read(buf[packetHeaderSize:])
		if err != nil {
			o.r = nil
			return
		}

		select {
		case replies <- buf[:packetHeaderSize+n]:
			// ok

		case <-o.reading:
			return
		}
	}
}
<commit_msg>service: origin: ignore zero-length writes<commit_after>package origin

import (
	"io"

	"github.com/tsavola/gate/service"
)

const (
	Name    = "origin"
	Version = 0

	packetHeaderSize = 8

	maxPacketSize = 0x10000 // TODO: move this elsewhere
)

type Factory struct {
	R io.Reader
	W io.Writer
}

func (f *Factory) Register(r *service.Registry) {
	service.Register(r, Name, Version, f)
}

func (f *Factory) New() service.Instance {
	return &origin{r: f.R, w: f.W}
}

var Default = new(Factory)

func Register(r *service.Registry) {
	Default.Register(r)
}

func CloneRegistryWith(r *service.Registry, origIn io.Reader, origOut io.Writer) *service.Registry {
	clone := service.Clone(r)
	(&Factory{R: origIn, W: origOut}).Register(clone)
	return clone
}

type origin struct {
	r io.Reader
	w io.Writer

	reading chan struct{}
}

func (o *origin) Handle(buf []byte, replies chan<- []byte) {
	if o.r != nil && o.reading == nil {
		o.reading = make(chan struct{})
		go o.readLoop(buf[6:8], replies)
	}

	if o.w != nil {
		if content := buf[packetHeaderSize:]; len(content) > 0 {
			if _, err := o.w.Write(content); err != nil {
				// assume that the error is EOF, broken pipe or such
				o.w = nil
			}
		}
	}
}

func (o *origin) Shutdown() {
	if o.reading != nil {
		close(o.reading)
	}
}

func (o *origin) readLoop(code []byte, replies chan<- []byte) {
	for {
		buf := make([]byte, maxPacketSize) // TODO: smaller buffer?
		copy(buf[6:8], code)

		n, err := o.r.Read(buf[packetHeaderSize:])
		if err != nil {
			o.r = nil
			return
		}

		select {
		case replies <- buf[:packetHeaderSize+n]:
			// ok

		case <-o.reading:
			return
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zanzibar

import (
	"context"
	"testing"

	"github.com/pborman/uuid"
	"github.com/stretchr/testify/assert"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest/observer"
)

func TestWithEndpointField(t *testing.T) {
	expected := "someEndpoint"
	ctx := WithEndpointField(context.TODO(), expected)

	ek := ctx.Value(endpointKey)
	endpoint, ok := ek.(string)

	assert.True(t, ok)
	assert.Equal(t, endpoint, expected)
}

func TestGetRequestEndpointFromCtx(t *testing.T) {
	expected := "someEndpoint"
	ctx := WithEndpointField(context.TODO(), expected)
	endpoint := GetRequestEndpointFromCtx(ctx)
	assert.Equal(t, expected, endpoint)

	expected = ""
	ctx = context.TODO()
	endpoint = GetRequestEndpointFromCtx(ctx)
	assert.Equal(t, expected, endpoint)
}

func TestWithEndpointRequestHeadersField(t *testing.T) {
	expected := map[string]string{"region": "san_francisco", "dc": "sjc1"}
	ctx := WithEndpointRequestHeadersField(context.TODO(), expected)
	rh := ctx.Value(endpointRequestHeader)
	requestHeaders, ok := rh.(map[string]string)

	assert.True(t, ok)
	assert.Equal(t, requestHeaders, expected)
}

func TestGetEndpointRequestHeadersFromCtx(t *testing.T) {
	expected := map[string]string{"region": "san_francisco", "dc": "sjc1"}
	headers := map[string]string{"region": "san_francisco", "dc": "sjc1"}
	ctx := WithEndpointRequestHeadersField(context.TODO(), headers)
	requestHeaders := GetEndpointRequestHeadersFromCtx(ctx)
	assert.Equal(t, expected, requestHeaders)

	expected = map[string]string{}
	ctx = context.TODO()
	requestHeaders = GetEndpointRequestHeadersFromCtx(ctx)
	assert.Equal(t, expected, requestHeaders)
}

func TestWithScopeTags(t *testing.T) {
	expected := map[string]string{"endpoint": "tincup", "handler": "exchange"}
	ctx := WithScopeTags(context.TODO(), expected)
	rs := ctx.Value(scopeTags)
	scopes, ok := rs.(map[string]string)

	assert.True(t, ok)
	assert.Equal(t, expected, scopes)
}

func TestGetScopeTagsFromCtx(t *testing.T) {
	expected := map[string]string{"endpoint": "tincup", "handler": "exchange"}
	scope := map[string]string{"endpoint": "tincup", "handler": "exchange"}
	ctx := WithScopeTags(context.TODO(), scope)
	scopes := GetScopeTagsFromCtx(ctx)
	assert.Equal(t, expected, scopes)

	expected = map[string]string{}
	ctx = context.TODO()
	scopes = GetScopeTagsFromCtx(ctx)
	assert.Equal(t, expected, scopes)
}

func TestWithRequestFields(t *testing.T) {
	uid := uuid.New()
	ctx := withRequestUUID(context.TODO(), uid)

	u := ctx.Value(requestUUIDKey)

	assert.NotNil(t, ctx)
	assert.Equal(t, uid, u)
}

func TestGetRequestUUIDFromCtx(t *testing.T) {
	uid := uuid.New()
	ctx := withRequestUUID(context.TODO(), uid)

	requestUUID := RequestUUIDFromCtx(ctx)

	assert.NotNil(t, ctx)
	assert.Equal(t, uid, requestUUID)

	// Test Default Scenario where no uuid exists in the context
	requestUUID = RequestUUIDFromCtx(context.TODO())
	assert.Equal(t, "", requestUUID)
}

func TestWithRoutingDelegate(t *testing.T) {
	expected := "somewhere"
	ctx := WithRoutingDelegate(context.TODO(), expected)
	rd := ctx.Value(routingDelegateKey)
	routingDelegate, ok := rd.(string)

	assert.True(t, ok)
	assert.Equal(t, routingDelegate, expected)
}

func TestGetRoutingDelegateFromCtx(t *testing.T) {
	expected := "somewhere"
	ctx := WithRoutingDelegate(context.TODO(), expected)
	rd := GetRoutingDelegateFromCtx(ctx)

	assert.Equal(t, expected, rd)
}

func TestWithShardKey(t *testing.T) {
	expected := "myshardkey"
	ctx := WithShardKey(context.TODO(), expected)
	sk := ctx.Value(shardKey)
	shardKey, ok := sk.(string)

	assert.True(t, ok)
	assert.Equal(t, shardKey, expected)
}

func TestGetShardKeyFromCtx(t *testing.T) {
	expected := "myshardkey"
	ctx := WithShardKey(context.TODO(), expected)
	sk := GetShardKeyFromCtx(ctx)

	assert.Equal(t, expected, sk)
}

func TestContextLogger(t *testing.T) {
	zapLoggerCore, logs := observer.New(zap.DebugLevel)
	zapLogger := zap.New(zapLoggerCore)
	contextLogger := NewContextLogger(zapLogger)
	ctx := context.Background()
	ctxWithField := WithLogFields(ctx, zap.String("ctxField", "ctxValue"))

	var logMessages []observer.LoggedEntry

	contextLogger.Debug(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.DebugLevel, logMessages[0].Level)
	assert.Equal(t, logMessages[0].Context[0].Key, "ctxField")
	assert.Equal(t, logMessages[0].Context[0].String, "ctxValue")
	assert.Equal(t, logMessages[0].Context[1].Key, "argField")
	assert.Equal(t, logMessages[0].Context[1].String, "argValue")

	contextLogger.Info(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.InfoLevel, logMessages[0].Level)

	contextLogger.WarnZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.WarnLevel, logMessages[0].Level)

	contextLogger.Error(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.ErrorLevel, logMessages[0].Level)
}

func TestContextLogger_DefaultZ(t *testing.T) {
	zapLoggerCore, logs := observer.New(zap.DebugLevel)
	zapLogger := zap.New(zapLoggerCore)
	contextLogger := NewContextLogger(zapLogger)
	ctx := context.Background()
	ctxWithField := WithLogFields(ctx, zap.String("ctxField", "ctxValue"))

	var logMessages []observer.LoggedEntry

	contextLogger.DebugZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.DebugLevel, logMessages[0].Level)
	assert.Equal(t, logMessages[0].Context[0].Key, "ctxField")
	assert.Equal(t, logMessages[0].Context[0].String, "ctxValue")
	assert.Equal(t, logMessages[0].Context[1].Key, "argField")
	assert.Equal(t, logMessages[0].Context[1].String, "argValue")

	contextLogger.InfoZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.InfoLevel, logMessages[0].Level)

	contextLogger.WarnZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.WarnLevel, logMessages[0].Level)

	contextLogger.ErrorZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.ErrorLevel, logMessages[0].Level)
}

func TestContextLogger_SkipZanzibarLogsZ(t *testing.T) {
	zapLoggerCore, logs := observer.New(zap.DebugLevel)
	zapLogger := zap.New(zapLoggerCore)
	contextLogger := NewContextLogger(zapLogger)
	contextLogger.SetSkipZanzibarLogs(true)
	ctx := context.Background()
	ctxWithField := WithLogFields(ctx, zap.String("ctxField", "ctxValue"))

	var logMessages []observer.LoggedEntry

	contextLogger.DebugZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 0)

	contextLogger.InfoZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 0)

	contextLogger.WarnZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 0)

	contextLogger.ErrorZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 0)
}

func TestContextLoggerPanic(t *testing.T) {
	defer func() {
		err := recover()
		assert.NotNil(t, err)
	}()

	zapNop := zap.NewNop()

	contextLogger := NewContextLogger(zapNop)
	ctx := context.Background()

	contextLogger.Panic(ctx, "msg", zap.String("argField", "argValue"))
}

func TestExtractScopeTag(t *testing.T) {
	headers := map[string]string{"x-uber-region-id": "san_francisco"}
	ctx := WithEndpointRequestHeadersField(context.TODO(), headers)
	contextScopeExtractors := []ContextScopeTagsExtractor{func(ctx context.Context) map[string]string {
		headers := GetEndpointRequestHeadersFromCtx(ctx)
		return map[string]string{"region-id": headers["x-uber-region-id"]}
	}}

	expected := map[string]string{"region-id": "san_francisco"}
	extractors := &ContextExtractors{
		ScopeTagsExtractors: contextScopeExtractors,
	}

	tags := extractors.ExtractScopeTags(ctx)
	assert.Equal(t, tags, expected)
}

func TestExtractLogField(t *testing.T) {
	headers := map[string]string{"x-uber-region-id": "san_francisco"}
	ctx := WithEndpointRequestHeadersField(context.TODO(), headers)
	contextLogFieldsExtractor := []ContextLogFieldsExtractor{func(ctx context.Context) []zap.Field {
		var fields []zap.Field
		headers := GetEndpointRequestHeadersFromCtx(ctx)
		fields = append(fields, zap.String("region-id", headers["x-uber-region-id"]))
		return fields
	}}

	var expected []zap.Field
	expected = append(expected, zap.String("region-id", "san_francisco"))
	extractors := &ContextExtractors{
		LogFieldsExtractors: contextLogFieldsExtractor,
	}
	fields := extractors.ExtractLogFields(ctx)
	assert.Equal(t, expected, fields)
}
<commit_msg>Improved tests<commit_after>// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zanzibar

import (
	"context"
	"testing"

	"github.com/pborman/uuid"
	"github.com/stretchr/testify/assert"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest/observer"
)

func TestWithEndpointField(t *testing.T) {
	expected := "someEndpoint"
	ctx := WithEndpointField(context.TODO(), expected)

	ek := ctx.Value(endpointKey)
	endpoint, ok := ek.(string)

	assert.True(t, ok)
	assert.Equal(t, endpoint, expected)
}

func TestGetRequestEndpointFromCtx(t *testing.T) {
	expected := "someEndpoint"
	ctx := WithEndpointField(context.TODO(), expected)
	endpoint := GetRequestEndpointFromCtx(ctx)
	assert.Equal(t, expected, endpoint)

	expected = ""
	ctx = context.TODO()
	endpoint = GetRequestEndpointFromCtx(ctx)
	assert.Equal(t, expected, endpoint)
}

func TestWithEndpointRequestHeadersField(t *testing.T) {
	expected := map[string]string{"region": "san_francisco", "dc": "sjc1"}
	ctx := WithEndpointRequestHeadersField(context.TODO(), expected)
	rh := ctx.Value(endpointRequestHeader)
	requestHeaders, ok := rh.(map[string]string)

	assert.True(t, ok)
	assert.Equal(t, requestHeaders, expected)
}

func TestGetEndpointRequestHeadersFromCtx(t *testing.T) {
	expected := map[string]string{"region": "san_francisco", "dc": "sjc1"}
	headers := map[string]string{"region": "san_francisco", "dc": "sjc1"}
	ctx := WithEndpointRequestHeadersField(context.TODO(), headers)
	requestHeaders := GetEndpointRequestHeadersFromCtx(ctx)
	assert.Equal(t, expected, requestHeaders)

	expected = map[string]string{}
	ctx = context.TODO()
	requestHeaders = GetEndpointRequestHeadersFromCtx(ctx)
	assert.Equal(t, expected, requestHeaders)
}

func TestWithScopeTags(t *testing.T) {
	expected := map[string]string{"endpoint": "tincup", "handler": "exchange"}
	ctx := WithScopeTags(context.TODO(), expected)
	rs := ctx.Value(scopeTags)
	scopes, ok := rs.(map[string]string)

	assert.True(t, ok)
	assert.Equal(t, expected, scopes)
}

func TestGetScopeTagsFromCtx(t *testing.T) {
	expected := map[string]string{"endpoint": "tincup", "handler": "exchange"}
	scope := map[string]string{"endpoint": "tincup", "handler": "exchange"}
	ctx := WithScopeTags(context.TODO(), scope)
	scopes := GetScopeTagsFromCtx(ctx)
	assert.Equal(t, expected, scopes)

	expected = map[string]string{}
	ctx = context.TODO()
	scopes = GetScopeTagsFromCtx(ctx)
	assert.Equal(t, expected, scopes)
}

func TestWithRequestFields(t *testing.T) {
	uid := uuid.New()
	ctx := withRequestUUID(context.TODO(), uid)

	u := ctx.Value(requestUUIDKey)

	assert.NotNil(t, ctx)
	assert.Equal(t, uid, u)
}

func TestGetRequestUUIDFromCtx(t *testing.T) {
	uid := uuid.New()
	ctx := withRequestUUID(context.TODO(), uid)

	requestUUID := RequestUUIDFromCtx(ctx)

	assert.NotNil(t, ctx)
	assert.Equal(t, uid, requestUUID)

	// Test Default Scenario where no uuid exists in the context
	requestUUID = RequestUUIDFromCtx(context.TODO())
	assert.Equal(t, "", requestUUID)
}

func TestWithRoutingDelegate(t *testing.T) {
	expected := "somewhere"
	ctx := WithRoutingDelegate(context.TODO(), expected)
	rd := ctx.Value(routingDelegateKey)
	routingDelegate, ok := rd.(string)

	assert.True(t, ok)
	assert.Equal(t, routingDelegate, expected)
}

func TestGetRoutingDelegateFromCtx(t *testing.T) {
	expected := "somewhere"
	ctx := WithRoutingDelegate(context.TODO(), expected)
	rd := GetRoutingDelegateFromCtx(ctx)

	assert.Equal(t, expected, rd)
}

func TestWithShardKey(t *testing.T) {
	expected := "myshardkey"
	ctx := WithShardKey(context.TODO(), expected)
	sk := ctx.Value(shardKey)
	shardKey, ok := sk.(string)

	assert.True(t, ok)
	assert.Equal(t, shardKey, expected)
}

func TestGetShardKeyFromCtx(t *testing.T) {
	expected := "myshardkey"
	ctx := WithShardKey(context.TODO(), expected)
	sk := GetShardKeyFromCtx(ctx)

	assert.Equal(t, expected, sk)
}

func TestContextLogger(t *testing.T) {
	zapLoggerCore, logs := observer.New(zap.DebugLevel)
	zapLogger := zap.New(zapLoggerCore)
	contextLogger := NewContextLogger(zapLogger)
	ctx := context.Background()
	ctxWithField := WithLogFields(ctx, zap.String("ctxField", "ctxValue"))

	var logMessages []observer.LoggedEntry

	contextLogger.Debug(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.DebugLevel, logMessages[0].Level)
	assert.Equal(t, logMessages[0].Context[0].Key, "ctxField")
	assert.Equal(t, logMessages[0].Context[0].String, "ctxValue")
	assert.Equal(t, logMessages[0].Context[1].Key, "argField")
	assert.Equal(t, logMessages[0].Context[1].String, "argValue")

	contextLogger.Info(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.InfoLevel, logMessages[0].Level)

	contextLogger.Warn(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.WarnLevel, logMessages[0].Level)

	contextLogger.Error(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.ErrorLevel, logMessages[0].Level)
}

func TestContextLogger_DefaultZ(t *testing.T) {
	zapLoggerCore, logs := observer.New(zap.DebugLevel)
	zapLogger := zap.New(zapLoggerCore)
	contextLogger := NewContextLogger(zapLogger)
	ctx := context.Background()
	ctxWithField := WithLogFields(ctx, zap.String("ctxField", "ctxValue"))

	var logMessages []observer.LoggedEntry

	contextLogger.DebugZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.DebugLevel, logMessages[0].Level)
	assert.Equal(t, logMessages[0].Context[0].Key, "ctxField")
	assert.Equal(t, logMessages[0].Context[0].String, "ctxValue")
	assert.Equal(t, logMessages[0].Context[1].Key, "argField")
	assert.Equal(t, logMessages[0].Context[1].String, "argValue")

	contextLogger.InfoZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.InfoLevel, logMessages[0].Level)

	contextLogger.WarnZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.WarnLevel, logMessages[0].Level)

	contextLogger.ErrorZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 1)
	assert.Equal(t, zap.ErrorLevel, logMessages[0].Level)
}

func TestContextLogger_SkipZanzibarLogsZ(t *testing.T) {
	zapLoggerCore, logs := observer.New(zap.DebugLevel)
	zapLogger := zap.New(zapLoggerCore)
	contextLogger := NewContextLogger(zapLogger)
	contextLogger.SetSkipZanzibarLogs(true)
	ctx := context.Background()
	ctxWithField := WithLogFields(ctx, zap.String("ctxField", "ctxValue"))

	var logMessages []observer.LoggedEntry

	contextLogger.DebugZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 0)

	contextLogger.InfoZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 0)

	contextLogger.WarnZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 0)

	contextLogger.ErrorZ(ctxWithField, "msg", zap.String("argField", "argValue"))
	logMessages = logs.TakeAll()
	assert.Len(t, logMessages, 0)
}

func TestContextLoggerPanic(t *testing.T) {
	defer func() {
		err := recover()
		assert.NotNil(t, err)
	}()

	zapNop := zap.NewNop()

	contextLogger := NewContextLogger(zapNop)
	ctx := context.Background()

	contextLogger.Panic(ctx, "msg", zap.String("argField", "argValue"))
}

func TestContextLoggerPanic_DefaultZ(t *testing.T) {
	defer func() {
		err := recover()
		assert.NotNil(t, err)
	}()

	zapNop := zap.NewNop()

	contextLogger := NewContextLogger(zapNop)
	ctx := context.Background()

	contextLogger.PanicZ(ctx, "msg", zap.String("argField", "argValue"))
}

func TestExtractScopeTag(t *testing.T) {
	headers := map[string]string{"x-uber-region-id": "san_francisco"}
	ctx := WithEndpointRequestHeadersField(context.TODO(), headers)
	contextScopeExtractors := []ContextScopeTagsExtractor{func(ctx context.Context) map[string]string {
		headers := GetEndpointRequestHeadersFromCtx(ctx)
		return map[string]string{"region-id": headers["x-uber-region-id"]}
	}}

	expected := map[string]string{"region-id": "san_francisco"}
	extractors := &ContextExtractors{
		ScopeTagsExtractors: contextScopeExtractors,
	}

	tags := extractors.ExtractScopeTags(ctx)
	assert.Equal(t, tags, expected)
}

func TestExtractLogField(t *testing.T) {
	headers := map[string]string{"x-uber-region-id": "san_francisco"}
	ctx := WithEndpointRequestHeadersField(context.TODO(), headers)
	contextLogFieldsExtractor := []ContextLogFieldsExtractor{func(ctx context.Context) []zap.Field {
		var fields []zap.Field
		headers := GetEndpointRequestHeadersFromCtx(ctx)
		fields = append(fields, zap.String("region-id", headers["x-uber-region-id"]))
		return fields
	}}

	var expected []zap.Field
	expected = append(expected, zap.String("region-id", "san_francisco"))
	extractors := &ContextExtractors{
		LogFieldsExtractors: contextLogFieldsExtractor,
	}
	fields := extractors.ExtractLogFields(ctx)
	assert.Equal(t, expected, fields)
}
<|endoftext|>


<commit_before>package daemon

import (
	"database/sql"
	"fmt"
	"path/filepath"

	"nimona.io/internal/net"
	"nimona.io/pkg/config"
	"nimona.io/pkg/context"
	"nimona.io/pkg/feedmanager"
	"nimona.io/pkg/hyperspace/resolver"
	"nimona.io/pkg/keystream"
	"nimona.io/pkg/network"
	"nimona.io/pkg/objectmanager"
	"nimona.io/pkg/objectstore"
	"nimona.io/pkg/peer"
	"nimona.io/pkg/preferences"
	"nimona.io/pkg/sqlobjectstore"
)

type (
	Daemon interface {
		Config() config.Config
		Preferences() preferences.Preferences
		Network() network.Network
		Resolver() resolver.Resolver
		ObjectStore() objectstore.Store
		ObjectManager() objectmanager.ObjectManager
		FeedManager() feedmanager.FeedManager
		KeyStreamManager() keystream.Manager
		// daemon specific methods
		Close()
	}
	daemon struct {
		config          config.Config
		preferences     preferences.Preferences
		configOptions   []config.Option
		network         network.Network
		resolver        resolver.Resolver
		objectstore     objectstore.Store
		objectmanager   objectmanager.ObjectManager
		feedmanager     feedmanager.FeedManager
		keystreamanager keystream.Manager
		// internal
		listener net.Listener
	}
	Option func(d *daemon) error
)

func New(ctx context.Context, opts ...Option) (Daemon, error) {
	d := &daemon{}

	// apply options
	for _, o := range opts {
		if err := o(d); err != nil {
			return nil, err
		}
	}

	// load config with given options
	cfg, err := config.New(d.configOptions...)
	if err != nil {
		return nil, fmt.Errorf("loading config: %w", err)
	}

	// construct new network
	ntw := network.New(
		ctx,
		network.WithPeerKey(cfg.Peer.PrivateKey),
	)

	if cfg.Peer.BindAddress != "" {
		// start listening
		lis, err := ntw.Listen(
			ctx,
			cfg.Peer.BindAddress,
			network.ListenOnLocalIPs,
			// network.ListenOnExternalPort,
		)
		if err != nil {
			return nil, fmt.Errorf("listening: %w", err)
		}
		d.listener = lis
	}

	// convert shorthands into connection infos
	bootstrapPeers := []*peer.ConnectionInfo{}
	for _, s := range cfg.Peer.Bootstraps {
		bootstrapPeer, err := s.GetConnectionInfo()
		if err != nil {
			return nil, fmt.Errorf("parsing bootstraps: %w", err)
		}
		bootstrapPeers = append(bootstrapPeers, bootstrapPeer)
	}

	// add bootstrap peers as relays
	ntw.RegisterRelays(bootstrapPeers...)

	// construct preferences db
	pdb, err := sql.Open("sqlite", filepath.Join(cfg.Path, "preferences.db"))
	if err != nil {
		return nil, fmt.Errorf("opening sql file for preferences: %w", err)
	}

	// construct preferences
	prf, err := preferences.NewSQLProvider(pdb)
	if err != nil {
		return nil, fmt.Errorf("constructing preferences provider: %w", err)
	}

	// construct object store
	db, err := sql.Open("sqlite", filepath.Join(cfg.Path, "nimona.db"))
	if err != nil {
		return nil, fmt.Errorf("opening sql file: %w", err)
	}

	str, err := sqlobjectstore.New(db)
	if err != nil {
		return nil, fmt.Errorf("starting sql store: %w", err)
	}

	// construct new resolver
	res := resolver.New(
		ctx,
		ntw,
		str,
		resolver.WithBoostrapPeers(bootstrapPeers...),
	)

	// construct manager
	man := objectmanager.New(
		ctx,
		ntw,
		res,
		str,
	)

	// construct feed manager
	fdm, err := feedmanager.New(
		ctx,
		ntw,
		res,
		str,
		man,
	)
	if err != nil {
		return nil, fmt.Errorf("constructing feed manager, %w", err)
	}

	// construct key stream manager
	ksm := keystream.NewKeyManager()

	d.config = *cfg
	d.preferences = prf
	d.network = ntw
	d.resolver = res
	d.objectstore = str
	d.objectmanager = man
	d.feedmanager = fdm
	d.keystreamanager = ksm

	return d, nil
}

func (d *daemon) Config() config.Config {
	return d.config
}

func (d *daemon) Preferences() preferences.Preferences {
	return d.preferences
}

func (d *daemon) Network() network.Network {
	return d.network
}

func (d *daemon) Resolver() resolver.Resolver {
	return d.resolver
}

func (d *daemon) ObjectStore() objectstore.Store {
	return d.objectstore
}

func (d *daemon) ObjectManager() objectmanager.ObjectManager {
	return d.objectmanager
}

func (d *daemon) FeedManager() feedmanager.FeedManager {
	return d.feedmanager
}

func (d *daemon) KeyStreamManager() keystream.Manager {
	return d.keystreamanager
}

func (d *daemon) Close() {
	if d.listener != nil {
		d.listener.Close() // nolint: errcheck
	}
}
<commit_msg>chore(daemon): fix keystream manager in daemon<commit_after>package daemon

import (
	"database/sql"
	"fmt"
	"path/filepath"

	"nimona.io/internal/net"
	"nimona.io/pkg/config"
	"nimona.io/pkg/context"
	"nimona.io/pkg/feedmanager"
	"nimona.io/pkg/hyperspace/resolver"
	"nimona.io/pkg/keystream"
	"nimona.io/pkg/network"
	"nimona.io/pkg/objectmanager"
	"nimona.io/pkg/objectstore"
	"nimona.io/pkg/peer"
	"nimona.io/pkg/preferences"
	"nimona.io/pkg/sqlobjectstore"
)

type (
	Daemon interface {
		Config() config.Config
		Preferences() preferences.Preferences
		Network() network.Network
		Resolver() resolver.Resolver
		ObjectStore() objectstore.Store
		ObjectManager() objectmanager.ObjectManager
		FeedManager() feedmanager.FeedManager
		KeyStreamManager() keystream.Manager
		// daemon specific methods
		Close()
	}
	daemon struct {
		config          config.Config
		preferences     preferences.Preferences
		configOptions   []config.Option
		network         network.Network
		resolver        resolver.Resolver
		objectstore     objectstore.Store
		objectmanager   objectmanager.ObjectManager
		feedmanager     feedmanager.FeedManager
		keystreamanager keystream.Manager
		// internal
		listener net.Listener
	}
	Option func(d *daemon) error
)

func New(ctx context.Context, opts ...Option) (Daemon, error) {
	d := &daemon{}

	// apply options
	for _, o := range opts {
		if err := o(d); err != nil {
			return nil, err
		}
	}

	// load config with given options
	cfg, err := config.New(d.configOptions...)
	if err != nil {
		return nil, fmt.Errorf("loading config: %w", err)
	}

	// construct new network
	ntw := network.New(
		ctx,
		network.WithPeerKey(cfg.Peer.PrivateKey),
	)

	if cfg.Peer.BindAddress != "" {
		// start listening
		lis, err := ntw.Listen(
			ctx,
			cfg.Peer.BindAddress,
			network.ListenOnLocalIPs,
			// network.ListenOnExternalPort,
		)
		if err != nil {
			return nil, fmt.Errorf("listening: %w", err)
		}
		d.listener = lis
	}

	// convert shorthands into connection infos
	bootstrapPeers := []*peer.ConnectionInfo{}
	for _, s := range cfg.Peer.Bootstraps {
		bootstrapPeer, err := s.GetConnectionInfo()
		if err != nil {
			return nil, fmt.Errorf("parsing bootstraps: %w", err)
		}
		bootstrapPeers = append(bootstrapPeers, bootstrapPeer)
	}

	// add bootstrap peers as relays
	ntw.RegisterRelays(bootstrapPeers...)

	// construct preferences db
	pdb, err := sql.Open("sqlite", filepath.Join(cfg.Path, "preferences.db"))
	if err != nil {
		return nil, fmt.Errorf("opening sql file for preferences: %w", err)
	}

	// construct preferences
	prf, err := preferences.NewSQLProvider(pdb)
	if err != nil {
		return nil, fmt.Errorf("constructing preferences provider: %w", err)
	}

	// construct object store
	db, err := sql.Open("sqlite", filepath.Join(cfg.Path, "nimona.db"))
	if err != nil {
		return nil, fmt.Errorf("opening sql file: %w", err)
	}

	str, err := sqlobjectstore.New(db)
	if err != nil {
		return nil, fmt.Errorf("starting sql store: %w", err)
	}

	// construct new resolver
	res := resolver.New(
		ctx,
		ntw,
		str,
		resolver.WithBoostrapPeers(bootstrapPeers...),
	)

	// construct manager
	man := objectmanager.New(
		ctx,
		ntw,
		res,
		str,
	)

	// construct feed manager
	fdm, err := feedmanager.New(
		ctx,
		ntw,
		res,
		str,
		man,
	)
	if err != nil {
		return nil, fmt.Errorf("constructing feed manager, %w", err)
	}

	// construct key stream manager
	ksm, err := keystream.NewKeyManager(
		ntw,
		str,
	)
	if err != nil {
		return nil, fmt.Errorf("constructing keystream manager, %w", err)
	}

	d.config = *cfg
	d.preferences = prf
	d.network = ntw
	d.resolver = res
	d.objectstore = str
	d.objectmanager = man
	d.feedmanager = fdm
	d.keystreamanager = ksm

	return d, nil
}

func (d *daemon) Config() config.Config {
	return d.config
}

func (d *daemon) Preferences() preferences.Preferences {
	return d.preferences
}

func (d *daemon) Network() network.Network {
	return d.network
}

func (d *daemon) Resolver() resolver.Resolver {
	return d.resolver
}

func (d *daemon) ObjectStore() objectstore.Store {
	return d.objectstore
}

func (d *daemon) ObjectManager() objectmanager.ObjectManager {
	return d.objectmanager
}

func (d *daemon) FeedManager() feedmanager.FeedManager {
	return d.feedmanager
}

func (d *daemon) KeyStreamManager() keystream.Manager {
	return d.keystreamanager
}

func (d *daemon) Close() {
	if d.listener != nil {
		d.listener.Close() // nolint: errcheck
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"

	"gopkg.in/go-playground/validator.v9"
)

var regexDBName = regexp.MustCompile(`^[a-z,A-Z,0-9,\.,\-,\_,\ ]+$`)
var regexPGTable = regexp.MustCompile(`^[a-z,A-Z,0-9,\.,\-,\_]+$`)
var regexSQLiteExpr = regexp.MustCompile(`^[a-z,A-Z,0-9,\.,\-,\_,\ ,\(,\),\',\",\%,\:]+$`)

// Custom validation function for SQLite database names
// At the moment it just allows alphanumeric and ".-_ " chars, though it should probably be extended to cover any
// valid file name
func checkDBName(fl validator.FieldLevel) bool {
	return regexDBName.MatchString(fl.Field().String())
}

// Custom validation function for PostgreSQL table names
// At the moment it just allows alphanumeric and ".-_" chars (may need to be expanded out at some point)
func checkPGTableName(fl validator.FieldLevel) bool {
	return regexPGTable.MatchString(fl.Field().String())
}

// Custom validation function for SQLite transformation expression
// eg this must pass:
//   printf('%s-%s-%s %s:00', substr(hour, 1, 4), substr(hour, 5, 2), substr(hour, 7, 2), substr(hour, 9, 2))
func checkSQLiteExpr(fl validator.FieldLevel) bool {
	return regexSQLiteExpr.MatchString(fl.Field().String())
}

// Checks a username against the list of reserved ones
func reservedUsernamesCheck(userName string) error {
	reserved := []string{"about", "admin", "blog", "download", "downloadcsv", "legal", "login", "logout", "mail",
		"news", "pref", "printer", "public", "reference", "register", "root", "star", "stars", "system",
		"table", "upload", "uploaddata", "vis"}
	for _, word := range reserved {
		if userName == word {
			return fmt.Errorf("That username is not available: %s\n", userName)
		}
	}

	return nil
}

// Validate the database name
func validateDB(dbName string) error {
	errs := validate.Var(dbName, "required,dbname,min=1,max=256") // 256 char limit seems reasonable
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided email address
func validateEmail(email string) error {
	errs := validate.Var(email, "required,email")
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided PostgreSQL table name
func validatePGTable(table string) error {
	// TODO: Improve this to work with all valid SQLite identifiers
	// TODO  Not seeing a definitive reference page for SQLite yet, so using the PostgreSQL one is
	// TODO  probably ok as a fallback:
	// TODO      https://www.postgresql.org/docs/current/static/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
	// TODO: Should we exclude SQLite internal tables too? (eg "sqlite_*" https://sqlite.org/lang_createtable.html)
	errs := validate.Var(table, "required,pgtable,max=63")
	if errs != nil {
		return errs
	}

	return nil
}

// Validate a user provided SQLite expression
func validateSQLiteexpr(user_expr string) error {
	errs := validate.Var(user_expr, "sqliteexpr,max=1024")
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided username
func validateUser(user string) error {
	errs := validate.Var(user, "required,alphanum,min=3,max=63")
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided user and database name
func validateUserDB(user string, db string) error {
	errs := validateUser(user)
	if errs != nil {
		return errs
	}

	errs = validateDB(db)
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided username and email address
func validateUserEmail(user string, email string) error {
	errs := validateUser(user)
	if errs != nil {
		return errs
	}

	errs = validateEmail(email)
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided user, database, and table name
func validateUserDBTable(user string, db string, table string) error {
	errs := validateUserDB(user, db)
	if errs != nil {
		return errs
	}

	errs = validatePGTable(table)
	if errs != nil {
		return errs
	}

	return nil
}
<commit_msg>Remove the now unused SQLite expression validator<commit_after>package main

import (
	"fmt"
	"regexp"

	"gopkg.in/go-playground/validator.v9"
)

var regexDBName = regexp.MustCompile(`^[a-z,A-Z,0-9,\.,\-,\_,\ ]+$`)
var regexPGTable = regexp.MustCompile(`^[a-z,A-Z,0-9,\.,\-,\_]+$`)

// Custom validation function for SQLite database names
// At the moment it just allows alphanumeric and ".-_ " chars, though it should probably be extended to cover any
// valid file name
func checkDBName(fl validator.FieldLevel) bool {
	return regexDBName.MatchString(fl.Field().String())
}

// Custom validation function for PostgreSQL table names
// At the moment it just allows alphanumeric and ".-_" chars (may need to be expanded out at some point)
func checkPGTableName(fl validator.FieldLevel) bool {
	return regexPGTable.MatchString(fl.Field().String())
}

// Checks a username against the list of reserved ones
func reservedUsernamesCheck(userName string) error {
	reserved := []string{"about", "admin", "blog", "download", "downloadcsv", "legal", "login", "logout", "mail",
		"news", "pref", "printer", "public", "reference", "register", "root", "star", "stars", "system",
		"table", "upload", "uploaddata", "vis"}
	for _, word := range reserved {
		if userName == word {
			return fmt.Errorf("That username is not available: %s\n", userName)
		}
	}

	return nil
}

// Validate the database name
func validateDB(dbName string) error {
	errs := validate.Var(dbName, "required,dbname,min=1,max=256") // 256 char limit seems reasonable
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided email address
func validateEmail(email string) error {
	errs := validate.Var(email, "required,email")
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided PostgreSQL table name
func validatePGTable(table string) error {
	// TODO: Improve this to work with all valid SQLite identifiers
	// TODO  Not seeing a definitive reference page for SQLite yet, so using the PostgreSQL one is
	// TODO  probably ok as a fallback:
	// TODO      https://www.postgresql.org/docs/current/static/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
	// TODO: Should we exclude SQLite internal tables too? (eg "sqlite_*" https://sqlite.org/lang_createtable.html)
	errs := validate.Var(table, "required,pgtable,max=63")
	if errs != nil {
		return errs
	}

	return nil
}

// Validate a user provided SQLite expression
func validateSQLiteexpr(user_expr string) error {
	errs := validate.Var(user_expr, "sqliteexpr,max=1024")
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided username
func validateUser(user string) error {
	errs := validate.Var(user, "required,alphanum,min=3,max=63")
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided user and database name
func validateUserDB(user string, db string) error {
	errs := validateUser(user)
	if errs != nil {
		return errs
	}

	errs = validateDB(db)
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided username and email address
func validateUserEmail(user string, email string) error {
	errs := validateUser(user)
	if errs != nil {
		return errs
	}

	errs = validateEmail(email)
	if errs != nil {
		return errs
	}

	return nil
}

// Validate the provided user, database, and table name
func validateUserDBTable(user string, db string, table string) error {
	errs := validateUserDB(user, db)
	if errs != nil {
		return errs
	}

	errs = validatePGTable(table)
	if errs != nil {
		return errs
	}

	return nil
}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"path"
	"runtime"
	"strings"
)

// Panic panicks
func Panic(msg string, prm ...interface{}) {
	CallerInfo(4)
	CallerInfo(3)
	CallerInfo(2)
	panic(fmt.Sprintf(msg, prm...))
}

// CallerInfo returns the file and line positions where an error occurred
//  idx -- use idx=2 to get the caller of Panic
func CallerInfo(idx int) {
	pc, file, line, ok := runtime.Caller(idx)
	if !ok {
		file, line = "?", 0
	}
	var fname string
	f := runtime.FuncForPC(pc)
	if f != nil {
		fname = f.Name()
	}
	if Verbose {
		// fmt.Printf("file = %s:%d\n", file, line)
		// fmt.Printf("func = %s\n", fname)
		fmt.Printf("file = %s:%d -- func = %s\n", file, line, fname)
	}
}

// GetErrMessage get error message if error
func GetErrMessage(err error) string {
	if err != nil {
		return fmt.Sprintf("Error is :'%s'", err.Error())
	}
	return "Notfound this error"
}

// CheckErr check error
func CheckErr(err error) {
	if err != nil {
		// perr(getErrMessage(err))
		panic(err)
	}
}

// FuncCallerInfo describes the information of callers of function
type FuncCallerInfo struct {
	PackageName string
	FileName    string
	FuncName    string
	Line        int
}

// RetrieveCallerInfo retrieve the information of callers of function
func RetrieveCallerInfo(idx int) *FuncCallerInfo {
	pc, file, line, _ := runtime.Caller(idx)
	_, fileName := path.Split(file)
	parts := strings.Split(runtime.FuncForPC(pc).Name(), ".")
	pl := len(parts)
	packageName := ""
	funcName := parts[pl-1]

	if parts[pl-2][0] == '(' {
		funcName = parts[pl-2] + "." + funcName
		packageName = strings.Join(parts[0:pl-2], ".")
	} else {
		packageName = strings.Join(parts[0:pl-1], ".")
	}

	return &FuncCallerInfo{
		PackageName: packageName,
		FileName:    fileName,
		FuncName:    funcName,
		Line:        line,
	}
}

// CallerName return the name of function calling
func CallerName(idx int) string {
	// pc, _, _, _ := runtime.Caller(idx) //idx = 0 self, 1 for caller, 2 for upper caller
	// msg := runtime.FuncForPC(pc).Name()
	ci := RetrieveCallerInfo(idx + 1)
	// if ci.fileName == "<autogenerated>" {
	// 	return Sf("package: %s:%d; func: %s",
	// 		LogColorString("debug", ci.packageName),
	// 		LogColorString("info", Sf("%d", ci.line)),
	// 		LogColorString("debug", ci.funcName),
	// 	)
	// }
	// return Sf("file: %s:%s; func: %s",
	// 	LogColorString("debug", ci.fileName),
	// 	LogColorString("info", Sf("%d", ci.line)),
	// 	LogColorString("debug", ci.funcName),
	// )
	// return Sf("file: %s:%d; func: %s", ci.FileName, ci.Line, ci.FuncName)
	return Sf("file: %s:%v; func: %s",
		ci.FileName,
		debugColor.Sprintf("%v", ci.Line),
		debugColor.Sprintf("%v", ci.FuncName),
	)
	// infos := make([]string, 0)
	// // if ci.FileName == "<autogenerated>" {
	// // 	infos = append(infos, HiYellowString("package: "))
	// // 	infos = append(infos, Color256String("%s", EncodeColor256(172, true), ci.PackageName))
	// // 	infos = append(infos, HiYellowString(":"))
	// // 	infos = append(infos, HiCyanString(Sf("%d", ci.Line)))
	// // 	infos = append(infos, HiYellowString("; func: "))
	// // 	infos = append(infos, Color256String("%s", EncodeColor256(175, true), ci.FuncName))
	// // } else {
	// infos = append(infos, HiYellowString("file: "))
	// infos = append(infos, Color256String("%s", EncodeColor256(172, true), ci.FileName))
	// infos = append(infos, HiYellowString(":"))
	// infos = append(infos, HiCyanString(Sf("%d", ci.Line)))
	// infos = append(infos, HiYellowString("; func: "))
	// infos = append(infos, Color256String("%s", EncodeColor256(175, true), ci.FuncName))
	// // }
	// return strings.Join(infos, "")
}

// Startfunc print the message at the start of function
func Startfunc(fid int) {
	Pfstart(Format[fid], CallerName(2))
}

// Stopfunc print the message at the start of function
func Stopfunc(fid int) {
	Pfstop(Format[fid], CallerName(2))
	PrintSepline(60)
}

// DebugPrintCaller print the name of function called and calling
func DebugPrintCaller() {
	// msg := Sf("▶ %q called by %q", CallerName(2), CallerName(3))
	// Glog.Debug(msg)
	if Glog.Printer.IsTerminal {
		Glog.Debugf("▶ [%s] called by [%s]",
			HiYellowString(CallerName(2)), YellowString(CallerName(3)))
	} else {
		Glog.Debugf("▶ [%s] called by [%s]", CallerName(2), CallerName(3))
	}
}
<commit_msg>remove dump codes<commit_after>package util

import (
	"fmt"
	"path"
	"runtime"
	"strings"
)

// Panic panicks
func Panic(msg string, prm ...interface{}) {
	CallerInfo(4)
	CallerInfo(3)
	CallerInfo(2)
	panic(fmt.Sprintf(msg, prm...))
}

// CallerInfo returns the file and line positions where an error occurred
//  idx -- use idx=2 to get the caller of Panic
func CallerInfo(idx int) {
	pc, file, line, ok := runtime.Caller(idx)
	if !ok {
		file, line = "?", 0
	}
	var fname string
	f := runtime.FuncForPC(pc)
	if f != nil {
		fname = f.Name()
	}
	if Verbose {
		// fmt.Printf("file = %s:%d\n", file, line)
		// fmt.Printf("func = %s\n", fname)
		fmt.Printf("file = %s:%d -- func = %s\n", file, line, fname)
	}
}

// GetErrMessage get error message if error
func GetErrMessage(err error) string {
	if err != nil {
		return fmt.Sprintf("Error is :'%s'", err.Error())
	}
	return "Notfound this error"
}

// CheckErr check error
func CheckErr(err error) {
	if err != nil {
		// perr(getErrMessage(err))
		panic(err)
	}
}

// FuncCallerInfo describes the information of callers of function
type FuncCallerInfo struct {
	PackageName string
	FileName    string
	FuncName    string
	Line        int
}

// RetrieveCallerInfo retrieve the information of callers of function
func RetrieveCallerInfo(idx int) *FuncCallerInfo {
	pc, file, line, _ := runtime.Caller(idx)
	_, fileName := path.Split(file)
	parts := strings.Split(runtime.FuncForPC(pc).Name(), ".")
	pl := len(parts)
	packageName := ""
	funcName := parts[pl-1]

	if parts[pl-2][0] == '(' {
		funcName = parts[pl-2] + "." + funcName
		packageName = strings.Join(parts[0:pl-2], ".")
	} else {
		packageName = strings.Join(parts[0:pl-1], ".")
	}

	return &FuncCallerInfo{
		PackageName: packageName,
		FileName:    fileName,
		FuncName:    funcName,
		Line:        line,
	}
}

// CallerName return the name of function calling
func CallerName(idx int) string {
	// pc, _, _, _ := runtime.Caller(idx) //idx = 0 self, 1 for caller, 2 for upper caller
	// msg := runtime.FuncForPC(pc).Name()
	ci := RetrieveCallerInfo(idx + 1)
	// if ci.fileName == "<autogenerated>" {
	// 	return Sf("package: %s:%d; func: %s",
	// 		LogColorString("debug", ci.packageName),
	// 		LogColorString("info", Sf("%d", ci.line)),
	// 		LogColorString("debug", ci.funcName),
	// 	)
	// }
	// return Sf("file: %s:%s; func: %s",
	// 	LogColorString("debug", ci.fileName),
	// 	LogColorString("info", Sf("%d", ci.line)),
	// 	LogColorString("debug", ci.funcName),
	// )
	// return Sf("file: %s:%d; func: %s", ci.FileName, ci.Line, ci.FuncName)
	return Sf("file: %s:%v; func: %s",
		ci.FileName,
		debugColor.Sprintf("%v", ci.Line),
		debugColor.Sprintf("%v", ci.FuncName),
	)
}

// Startfunc print the message at the start of function
func Startfunc(fid int) {
	Pfstart(Format[fid], CallerName(2))
}

// Stopfunc print the message at the start of function
func Stopfunc(fid int) {
	Pfstop(Format[fid], CallerName(2))
	PrintSepline(60)
}

// DebugPrintCaller print the name of function called and calling
func DebugPrintCaller() {
	// msg := Sf("▶ %q called by %q", CallerName(2), CallerName(3))
	// Glog.Debug(msg)
	if Glog.Printer.IsTerminal {
		Glog.Debugf("▶ [%s] called by [%s]",
			HiYellowString(CallerName(2)), YellowString(CallerName(3)))
	} else {
		Glog.Debugf("▶ [%s] called by [%s]", CallerName(2), CallerName(3))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"github.com/anacrolix/utp"
)

var (
	flClientMode = flag.Bool("c", false, "client mode")
	flServerMode = flag.Bool("s", false, "server mode")
	flHost       = flag.String("h", "127.0.0.1", "host")
	flPort       = flag.Int("p", 6001, "port")
	flLen        = flag.Int("l", 1400, "length of data")
	flThreads    = flag.Int("t", 1, "threads")
	flDuration   = flag.Duration("d", time.Second*10, "duration")
)

func main() {
	log.Printf("UTP Benchmark Tool by Artem Andreenko (miolini@gmail.com)")
	flag.Parse()
	if !*flClientMode && !*flServerMode {
		flag.PrintDefaults()
		return
	}
	ts := time.Now()
	wg := sync.WaitGroup{}
	if *flServerMode {
		wg.Add(1)
		go server(&wg, *flHost, *flPort)
	}
	if *flClientMode {
		wg.Add(*flThreads)
		chStat := make(chan int, 100)
		go stat(chStat)
		for i := 0; i < *flThreads; i++ {
			go client(&wg, *flHost, *flPort, *flLen, *flDuration, chStat)
		}
	}
	wg.Wait()
	log.Printf("time takes %.2fsec", time.Since(ts).Seconds())
}

func stat(chStat chan int) {
	t := time.NewTicker(time.Second)
	counter := 0
	for {
		select {
		case n := <-chStat:
			counter += n
		case <-t.C:
			log.Printf("speed %.3f mbit/sec", float64(counter*8)/1024/1024)
			counter = 0
		}
	}
}

func server(wg *sync.WaitGroup, host string, port int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("error: %s", r)
		}
		wg.Done()
	}()
	log.Printf("server listen %s:%d", host, port)
	s, err := utp.NewSocket("udp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		panic(err)
	}
	defer s.Close()
	for {
		conn, err := s.Accept()
		if err != nil {
			panic(err)
		}
		wg.Add(1)
		go readConn(conn)
	}
}

func readConn(conn net.Conn) {
	defer conn.Close()
	defer log.Printf("client %s disconnected", conn.RemoteAddr().String())
	log.Printf("client %s connected", conn.RemoteAddr().String())
	buf := make([]byte, 4096)
	for {
		_, err := conn.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("err: %s", err)
			}
			break
		}
	}
}

func client(wg *sync.WaitGroup, host string, port, len int, duration time.Duration, chStat chan int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("error: %s", r)
		}
		log.Printf("disconnected")
		wg.Done()
	}()
	log.Printf("connecting to %s:%d, len %d, duration %s", host, port, len, duration.String())
	conn, err := utp.DialTimeout(fmt.Sprintf("%s:%d", host, port), time.Second)
	if err != nil {
		panic(err)
	}
	defer conn.Close()
	log.Printf("connected")
	buf := bytes.Repeat([]byte("H"), len)
	ts := time.Now()
	for time.Since(ts) < duration {
		n, err := conn.Write(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}
		chStat <- n
	}
}
<commit_msg>duration for stats<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"github.com/anacrolix/utp"
)

var (
	flClientMode   = flag.Bool("c", false, "client mode")
	flServerMode   = flag.Bool("s", false, "server mode")
	flHost         = flag.String("h", "127.0.0.1", "host")
	flPort         = flag.Int("p", 6001, "port")
	flLen          = flag.Int("l", 1400, "length of data")
	flThreads      = flag.Int("t", 1, "threads")
	flDuration     = flag.Duration("d", time.Second*10, "duration")
	flDurationStat = flag.Duration("ds", time.Second*5, "duration for stats")
)

func main() {
	log.Printf("UTP Benchmark Tool by Artem Andreenko (miolini@gmail.com)")
	flag.Parse()
	if !*flClientMode && !*flServerMode {
		flag.PrintDefaults()
		return
	}
	ts := time.Now()
	wg := sync.WaitGroup{}
	if *flServerMode {
		wg.Add(1)
		go server(&wg, *flHost, *flPort)
	}
	if *flClientMode {
		wg.Add(*flThreads)
		chStat := make(chan int, 100)
		go stat(chStat, *flDurationStat)
		for i := 0; i < *flThreads; i++ {
			go client(&wg, *flHost, *flPort, *flLen, *flDuration, chStat)
		}
	}
	wg.Wait()
	log.Printf("time takes %.2fsec", time.Since(ts).Seconds())
}

func stat(chStat chan int, duration time.Duration) {
	t := time.NewTicker(time.Second)
	counter := 0
	for {
		select {
		case n := <-chStat:
			counter += n
		case <-t.C:
			log.Printf("speed %.3f mbit/sec", float64(counter*8)/duration.Seconds()/1024/1024)
			counter = 0
		}
	}
}

func server(wg *sync.WaitGroup, host string, port int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("error: %s", r)
		}
		wg.Done()
	}()
	log.Printf("server listen %s:%d", host, port)
	s, err := utp.NewSocket("udp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		panic(err)
	}
	defer s.Close()
	for {
		conn, err := s.Accept()
		if err != nil {
			panic(err)
		}
		wg.Add(1)
		go readConn(conn)
	}
}

func readConn(conn net.Conn) {
	defer conn.Close()
	defer log.Printf("client %s disconnected", conn.RemoteAddr().String())
	log.Printf("client %s connected", conn.RemoteAddr().String())
	buf := make([]byte, 4096)
	for {
		_, err := conn.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("err: %s", err)
			}
			break
		}
	}
}

func client(wg *sync.WaitGroup, host string, port, len int, duration time.Duration, chStat chan int) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("error: %s", r)
		}
		log.Printf("disconnected")
		wg.Done()
	}()
	log.Printf("connecting to %s:%d, len %d, duration %s", host, port, len, duration.String())
	conn, err := utp.DialTimeout(fmt.Sprintf("%s:%d", host, port), time.Second)
	if err != nil {
		panic(err)
	}
	defer conn.Close()
	log.Printf("connected")
	buf := bytes.Repeat([]byte("H"), len)
	ts := time.Now()
	for time.Since(ts) < duration {
		n, err := conn.Write(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}
		chStat <- n
	}
}
<|endoftext|>


<commit_before>package xls

import (
	"bytes"
	"encoding/binary"
	"io"
	"os"
	"unicode/utf16"
	"golang.org/x/text/encoding/charmap"
)

//xls workbook type
type WorkBook struct {
	Is5ver   bool
	Type     uint16
	Codepage uint16
	Xfs      []st_xf_data
	Fonts    []Font
	Formats  map[uint16]*Format
	//All the sheets from the workbook
	sheets         []*WorkSheet
	Author         string
	rs             io.ReadSeeker
	sst            []string
	continue_utf16 uint16
	continue_rich  uint16
	continue_apsb  uint32
	dateMode       uint16
}

//read workbook from ole2 file
func newWorkBookFromOle2(rs io.ReadSeeker) *WorkBook {
	wb := new(WorkBook)
	wb.Formats = make(map[uint16]*Format)
	// wb.bts = bts
	wb.rs = rs
	wb.sheets = make([]*WorkSheet, 0)
	wb.Parse(rs)
	return wb
}

func (w *WorkBook) Parse(buf io.ReadSeeker) {
	b := new(bof)
	bof_pre := new(bof)
	// buf := bytes.NewReader(bts)
	offset := 0
	for {
		if err := binary.Read(buf, binary.LittleEndian, b); err == nil {
			bof_pre, b, offset = w.parseBof(buf, b, bof_pre, offset)
		} else {
			break
		}
	}
}

func (w *WorkBook) addXf(xf st_xf_data) {
	w.Xfs = append(w.Xfs, xf)
}

func (w *WorkBook) addFont(font *FontInfo, buf io.ReadSeeker) {
	name, _ := w.get_string(buf, uint16(font.NameB))
	w.Fonts = append(w.Fonts, Font{Info: font, Name: name})
}

func (w *WorkBook) addFormat(format *Format) {
	if w.Formats == nil {
		os.Exit(1)
	}
	w.Formats[format.Head.Index] = format
}

func (wb *WorkBook) parseBof(buf io.ReadSeeker, b *bof, pre *bof, offset_pre int) (after *bof, after_using *bof, offset int) {
	after = b
	after_using = pre
	var bts = make([]byte, b.Size)
	binary.Read(buf, binary.LittleEndian, bts)
	buf_item := bytes.NewReader(bts)
	switch b.Id {
	case 0x809:
		bif := new(biffHeader)
		binary.Read(buf_item, binary.LittleEndian, bif)
		if bif.Ver != 0x600 {
			wb.Is5ver = true
		}
		wb.Type = bif.Type
	case 0x042: // CODEPAGE
		binary.Read(buf_item, binary.LittleEndian, &wb.Codepage)
	case 0x3c: // CONTINUE
		// step back if previous element not yet completed
		if wb.continue_utf16 > 0 {
			offset_pre--
		}

		if pre.Id == 0xfc {
			var size uint16
			var err error
			if wb.continue_utf16 >= 1 {
				size = wb.continue_utf16
				wb.continue_utf16 = 0
			} else {
				err = binary.Read(buf_item, binary.LittleEndian, &size)
			}
			for err == nil && offset_pre < len(wb.sst) {
				var str string
				if size > 0 {
					str, err = wb.get_string(buf_item, size)
					wb.sst[offset_pre] = wb.sst[offset_pre] + str
				}

				if err == io.EOF {
					break
				}

				offset_pre++
				err = binary.Read(buf_item, binary.LittleEndian, &size)
			}
		}
		offset = offset_pre
		after = pre
		after_using = b
	case 0xfc: // SST
		info := new(SstInfo)
		binary.Read(buf_item, binary.LittleEndian, info)
		wb.sst = make([]string, info.Count)
		var size uint16
		var i = 0
		// dont forget to initialize offset
		offset = 0
		for ; i < int(info.Count); i++ {
			var err error
			if err = binary.Read(buf_item, binary.LittleEndian, &size); err == nil {
				var str string
				str, err = wb.get_string(buf_item, size)
				wb.sst[i] = wb.sst[i] + str
			}

			if err == io.EOF {
				break
			}
		}
		offset = i
	case 0x85: // bOUNDSHEET
		var bs = new(boundsheet)
		binary.Read(buf_item, binary.LittleEndian, bs)
		// different for BIFF5 and BIFF8
		wb.addSheet(bs, buf_item)
	case 0x0e0: // XF
		if wb.Is5ver {
			xf := new(Xf5)
			binary.Read(buf_item, binary.LittleEndian, xf)
			wb.addXf(xf)
		} else {
			xf := new(Xf8)
			binary.Read(buf_item, binary.LittleEndian, xf)
			wb.addXf(xf)
		}
	case 0x031: // FONT
		f := new(FontInfo)
		binary.Read(buf_item, binary.LittleEndian, f)
		wb.addFont(f, buf_item)
	case 0x41E: //FORMAT
		font := new(Format)
		binary.Read(buf_item, binary.LittleEndian, &font.Head)
		font.str, _ = wb.get_string(buf_item, font.Head.Size)
		wb.addFormat(font)
	case 0x22: //DATEMODE
		binary.Read(buf_item, binary.LittleEndian, &wb.dateMode)
	}
	return
}
func decodeWindows1251(enc []byte) string {
    dec := charmap.Windows1251.NewDecoder()
    out, _ := dec.Bytes(enc)
    return string(out)
}
func (w *WorkBook) get_string(buf io.ReadSeeker, size uint16) (res string, err error) {
	if w.Is5ver {
		var bts = make([]byte, size)
		_, err = buf.Read(bts)
		res = decodeWindows1251(bts)
		//res = string(bts)
	} else {
		var richtext_num = uint16(0)
		var phonetic_size = uint32(0)
		var flag byte
		err = binary.Read(buf, binary.LittleEndian, &flag)
		if flag&0x8 != 0 {
			err = binary.Read(buf, binary.LittleEndian, &richtext_num)
		} else if w.continue_rich > 0 {
			richtext_num = w.continue_rich
			w.continue_rich = 0
		}
		if flag&0x4 != 0 {
			err = binary.Read(buf, binary.LittleEndian, &phonetic_size)
		} else if w.continue_apsb > 0 {
			phonetic_size = w.continue_apsb
			w.continue_apsb = 0
		}
		if flag&0x1 != 0 {
			var bts = make([]uint16, size)
			var i = uint16(0)
			// we need local err here
			var err error
			for ; i < size && err == nil; i++ {
				err = binary.Read(buf, binary.LittleEndian, &bts[i])
			}

			// when eof found, we dont want to append last element
			var runes []rune
			if err == io.EOF {
				runes = utf16.Decode(bts[:i-1])
			} else {
				runes = utf16.Decode(bts[:i])
			}

			res = string(runes)
			if i < size {
				w.continue_utf16 = size - i + 1
			}
		} else {
			var bts = make([]byte, size)
			var n int
			n, err = buf.Read(bts)
			if uint16(n) < size {
				w.continue_utf16 = size - uint16(n)
				err = io.EOF
			}

			var bts1 = make([]uint16, n)
			for k, v := range bts[:n] {
				bts1[k] = uint16(v)
			}
			runes := utf16.Decode(bts1)
			res = string(runes)
		}
		if richtext_num > 0 {
			var bts []byte
			var seek_size int64
			if w.Is5ver {
				seek_size = int64(2 * richtext_num)
			} else {
				seek_size = int64(4 * richtext_num)
			}
			bts = make([]byte, seek_size)
			err = binary.Read(buf, binary.LittleEndian, bts)
			if err == io.EOF {
				w.continue_rich = richtext_num
			}

			// err = binary.Read(buf, binary.LittleEndian, bts)
		}
		if phonetic_size > 0 {
			var bts []byte
			bts = make([]byte, phonetic_size)
			err = binary.Read(buf, binary.LittleEndian, bts)
			if err == io.EOF {
				w.continue_apsb = phonetic_size
			}
		}
	}
	return
}

func (w *WorkBook) addSheet(sheet *boundsheet, buf io.ReadSeeker) {
	name, _ := w.get_string(buf, uint16(sheet.Name))
	w.sheets = append(w.sheets, &WorkSheet{bs: sheet, Name: name, wb: w})
}

//reading a sheet from the compress file to memory, you should call this before you try to get anything from sheet
func (w *WorkBook) prepareSheet(sheet *WorkSheet) {
	w.rs.Seek(int64(sheet.bs.Filepos), 0)
	sheet.parse(w.rs)
}

//Get one sheet by its number
func (w *WorkBook) GetSheet(num int) *WorkSheet {
	if num < len(w.sheets) {
		s := w.sheets[num]
		if !s.parsed {
			w.prepareSheet(s)
		}
		return s
	} else {
		return nil
	}
}

//Get the number of all sheets, look into example
func (w *WorkBook) NumSheets() int {
	return len(w.sheets)
}

//helper function to read all cells from file
//Notice: the max value is the limit of the max capacity of lines.
//Warning: the helper function will need big memeory if file is large.
func (w *WorkBook) ReadAllCells(max int) (res [][]string) {
	res = make([][]string, 0)
	for _, sheet := range w.sheets {
		if len(res) < max {
			max = max - len(res)
			w.prepareSheet(sheet)
			if sheet.MaxRow != 0 {
				leng := int(sheet.MaxRow) + 1
				if max < leng {
					leng = max
				}
				temp := make([][]string, leng)
				for k, row := range sheet.rows {
					data := make([]string, 0)
					if len(row.cols) > 0 {
						for _, col := range row.cols {
							if uint16(len(data)) <= col.LastCol() {
								data = append(data, make([]string, col.LastCol()-uint16(len(data))+1)...)
							}
							str := col.String(w)

							for i := uint16(0); i < col.LastCol()-col.FirstCol()+1; i++ {
								data[col.FirstCol()+i] = str[i]
							}
						}
						if leng > int(k) {
							temp[k] = data
						}
					}
				}
				res = append(res, temp...)
			}
		}
	}
	return
}
<commit_msg>FIX: issue #47<commit_after>package xls

import (
	"bytes"
	"encoding/binary"
	"golang.org/x/text/encoding/charmap"
	"io"
	"os"
	"unicode/utf16"
)

//xls workbook type
type WorkBook struct {
	Is5ver   bool
	Type     uint16
	Codepage uint16
	Xfs      []st_xf_data
	Fonts    []Font
	Formats  map[uint16]*Format
	//All the sheets from the workbook
	sheets         []*WorkSheet
	Author         string
	rs             io.ReadSeeker
	sst            []string
	continue_utf16 uint16
	continue_rich  uint16
	continue_apsb  uint32
	dateMode       uint16
}

//read workbook from ole2 file
func newWorkBookFromOle2(rs io.ReadSeeker) *WorkBook {
	wb := new(WorkBook)
	wb.Formats = make(map[uint16]*Format)
	// wb.bts = bts
	wb.rs = rs
	wb.sheets = make([]*WorkSheet, 0)
	wb.Parse(rs)
	return wb
}

func (w *WorkBook) Parse(buf io.ReadSeeker) {
	b := new(bof)
	bof_pre := new(bof)
	// buf := bytes.NewReader(bts)
	offset := 0
	for {
		if err := binary.Read(buf, binary.LittleEndian, b); err == nil {
			bof_pre, b, offset = w.parseBof(buf, b, bof_pre, offset)
		} else {
			break
		}
	}
}

func (w *WorkBook) addXf(xf st_xf_data) {
	w.Xfs = append(w.Xfs, xf)
}

func (w *WorkBook) addFont(font *FontInfo, buf io.ReadSeeker) {
	name, _ := w.get_string(buf, uint16(font.NameB))
	w.Fonts = append(w.Fonts, Font{Info: font, Name: name})
}

func (w *WorkBook) addFormat(format *Format) {
	if w.Formats == nil {
		os.Exit(1)
	}
	w.Formats[format.Head.Index] = format
}

func (wb *WorkBook) parseBof(buf io.ReadSeeker, b *bof, pre *bof, offset_pre int) (after *bof, after_using *bof, offset int) {
	after = b
	after_using = pre
	var bts = make([]byte, b.Size)
	binary.Read(buf, binary.LittleEndian, bts)
	buf_item := bytes.NewReader(bts)
	switch b.Id {
	case 0x809:
		bif := new(biffHeader)
		binary.Read(buf_item, binary.LittleEndian, bif)
		if bif.Ver != 0x600 {
			wb.Is5ver = true
		}
		wb.Type = bif.Type
	case 0x042: // CODEPAGE
		binary.Read(buf_item, binary.LittleEndian, &wb.Codepage)
	case 0x3c: // CONTINUE
		if pre.Id == 0xfc {
			var size uint16
			var err error
			if wb.continue_utf16 >= 1 {
				size = wb.continue_utf16
				wb.continue_utf16 = 0
			} else {
				err = binary.Read(buf_item, binary.LittleEndian, &size)
			}
			for err == nil && offset_pre < len(wb.sst) {
				var str string
				str, err = wb.get_string(buf_item, size)
				wb.sst[offset_pre] = wb.sst[offset_pre] + str

				if err == io.EOF {
					break
				}

				offset_pre++
				err = binary.Read(buf_item, binary.LittleEndian, &size)
			}
		}
		offset = offset_pre
		after = pre
		after_using = b
	case 0xfc: // SST
		info := new(SstInfo)
		binary.Read(buf_item, binary.LittleEndian, info)
		wb.sst = make([]string, info.Count)
		var size uint16
		var i = 0
		// dont forget to initialize offset
		offset = 0
		for ; i < int(info.Count); i++ {
			var err error
			err = binary.Read(buf_item, binary.LittleEndian, &size)
			if err == nil {
				var str string
				str, err = wb.get_string(buf_item, size)
				wb.sst[i] = wb.sst[i] + str
			}

			if err == io.EOF {
				break
			}
		}
		offset = i
	case 0x85: // boundsheet
		var bs = new(boundsheet)
		binary.Read(buf_item, binary.LittleEndian, bs)
		// different for BIFF5 and BIFF8
		wb.addSheet(bs, buf_item)
	case 0x0e0: // XF
		if wb.Is5ver {
			xf := new(Xf5)
			binary.Read(buf_item, binary.LittleEndian, xf)
			wb.addXf(xf)
		} else {
			xf := new(Xf8)
			binary.Read(buf_item, binary.LittleEndian, xf)
			wb.addXf(xf)
		}
	case 0x031: // FONT
		f := new(FontInfo)
		binary.Read(buf_item, binary.LittleEndian, f)
		wb.addFont(f, buf_item)
	case 0x41E: //FORMAT
		font := new(Format)
		binary.Read(buf_item, binary.LittleEndian, &font.Head)
		font.str, _ = wb.get_string(buf_item, font.Head.Size)
		wb.addFormat(font)
	case 0x22: //DATEMODE
		binary.Read(buf_item, binary.LittleEndian, &wb.dateMode)
	}
	return
}
func decodeWindows1251(enc []byte) string {
	dec := charmap.Windows1251.NewDecoder()
	out, _ := dec.Bytes(enc)
	return string(out)
}
func (w *WorkBook) get_string(buf io.ReadSeeker, size uint16) (res string, err error) {
	if w.Is5ver {
		var bts = make([]byte, size)
		_, err = buf.Read(bts)
		res = decodeWindows1251(bts)
		//res = string(bts)
	} else {
		var richtext_num = uint16(0)
		var phonetic_size = uint32(0)
		var flag byte
		err = binary.Read(buf, binary.LittleEndian, &flag)
		if flag&0x8 != 0 {
			err = binary.Read(buf, binary.LittleEndian, &richtext_num)
		} else if w.continue_rich > 0 {
			richtext_num = w.continue_rich
			w.continue_rich = 0
		}
		if flag&0x4 != 0 {
			err = binary.Read(buf, binary.LittleEndian, &phonetic_size)
		} else if w.continue_apsb > 0 {
			phonetic_size = w.continue_apsb
			w.continue_apsb = 0
		}
		if flag&0x1 != 0 {
			var bts = make([]uint16, size)
			var i = uint16(0)
			for ; i < size && err == nil; i++ {
				err = binary.Read(buf, binary.LittleEndian, &bts[i])
			}

			// when eof found, we dont want to append last element
			var runes []rune
			if err == io.EOF {
				i = i - 1
			}
			runes = utf16.Decode(bts[:i])

			res = string(runes)
			if i < size {
				w.continue_utf16 = size - i
			}

		} else {
			var bts = make([]byte, size)
			var n int
			n, err = buf.Read(bts)
			if uint16(n) < size {
				w.continue_utf16 = size - uint16(n)
				err = io.EOF
			}

			var bts1 = make([]uint16, n)
			for k, v := range bts[:n] {
				bts1[k] = uint16(v)
			}
			runes := utf16.Decode(bts1)
			res = string(runes)
		}
		if richtext_num > 0 {
			var bts []byte
			var seek_size int64
			if w.Is5ver {
				seek_size = int64(2 * richtext_num)
			} else {
				seek_size = int64(4 * richtext_num)
			}
			bts = make([]byte, seek_size)
			err = binary.Read(buf, binary.LittleEndian, bts)
			if err == io.EOF {
				w.continue_rich = richtext_num
			}

			// err = binary.Read(buf, binary.LittleEndian, bts)
		}
		if phonetic_size > 0 {
			var bts []byte
			bts = make([]byte, phonetic_size)
			err = binary.Read(buf, binary.LittleEndian, bts)
			if err == io.EOF {
				w.continue_apsb = phonetic_size
			}
		}
	}
	return
}

func (w *WorkBook) addSheet(sheet *boundsheet, buf io.ReadSeeker) {
	name, _ := w.get_string(buf, uint16(sheet.Name))
	w.sheets = append(w.sheets, &WorkSheet{bs: sheet, Name: name, wb: w})
}

//reading a sheet from the compress file to memory, you should call this before you try to get anything from sheet
func (w *WorkBook) prepareSheet(sheet *WorkSheet) {
	w.rs.Seek(int64(sheet.bs.Filepos), 0)
	sheet.parse(w.rs)
}

//Get one sheet by its number
func (w *WorkBook) GetSheet(num int) *WorkSheet {
	if num < len(w.sheets) {
		s := w.sheets[num]
		if !s.parsed {
			w.prepareSheet(s)
		}
		return s
	} else {
		return nil
	}
}

//Get the number of all sheets, look into example
func (w *WorkBook) NumSheets() int {
	return len(w.sheets)
}

//helper function to read all cells from file
//Notice: the max value is the limit of the max capacity of lines.
//Warning: the helper function will need big memeory if file is large.
func (w *WorkBook) ReadAllCells(max int) (res [][]string) {
	res = make([][]string, 0)
	for _, sheet := range w.sheets {
		if len(res) < max {
			max = max - len(res)
			w.prepareSheet(sheet)
			if sheet.MaxRow != 0 {
				leng := int(sheet.MaxRow) + 1
				if max < leng {
					leng = max
				}
				temp := make([][]string, leng)
				for k, row := range sheet.rows {
					data := make([]string, 0)
					if len(row.cols) > 0 {
						for _, col := range row.cols {
							if uint16(len(data)) <= col.LastCol() {
								data = append(data, make([]string, col.LastCol()-uint16(len(data))+1)...)
							}
							str := col.String(w)

							for i := uint16(0); i < col.LastCol()-col.FirstCol()+1; i++ {
								data[col.FirstCol()+i] = str[i]
							}
						}
						if leng > int(k) {
							temp[k] = data
						}
					}
				}
				res = append(res, temp...)
			}
		}
	}
	return
}
<|endoftext|>