<commit_before>package main

import (
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"log"
	"net/rpc"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"runtime/debug"

	"github.com/mdempsky/gocode/internal/cache"
	"github.com/mdempsky/gocode/internal/suggest"
)

func doClient() {
	// Client is a short-lived program.
	// Disable GC to make it faster
	debug.SetGCPercent(-1)

	if *g_debug {
		start := time.Now()
		defer func() {
			elapsed := time.Since(start)
			log.Printf("Elapsed duration: %v\n", elapsed)
		}()
	}

	var command string
	if flag.NArg() > 0 {
		command = flag.Arg(0)
		switch command {
		case "autocomplete", "exit":
			// these are valid commands
		case "close":
			// "close" is an alias for "exit"
			command = "exit"
		default:
			fmt.Printf("gocode: unknown subcommand: %q\nRun 'gocode -help' for usage.\n", command)
			os.Exit(2)
		}
	}

	// client
	var client *rpc.Client
	if *g_sock != "none" {
		addr := *g_addr
		if *g_sock == "unix" {
			addr = getSocketPath()
		}

		var err error
		client, err = rpc.Dial(*g_sock, addr)
		if err != nil {
			if command == "exit" {
				log.Fatal(err)
			}

			if *g_sock == "unix" {
				_ = os.Remove(addr)
			}
			err = tryStartServer()
			if err != nil {
				log.Fatalf("Failed to start server: %s\n", err)
			}
			client, err = tryToConnect(*g_sock, addr)
			if err != nil {
				log.Fatalf("Failed to connect to %q: %s\n", addr, err)
			}
		}
		defer client.Close()
	}

	switch command {
	case "autocomplete":
		cmdAutoComplete(client)
	case "exit":
		cmdExit(client)
	}
}

func tryStartServer() error {
	path := get_executable_filename()
	args := []string{os.Args[0], "-s", "-sock", *g_sock, "-addr", *g_addr}
	cwd, _ := os.Getwd()

	var err error
	stdin, err := os.Open(os.DevNull)
	if err != nil {
		return err
	}
	stdout, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
	if err != nil {
		return err
	}
	stderr, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
	if err != nil {
		return err
	}

	procattr := os.ProcAttr{Dir: cwd, Env: os.Environ(), Files: []*os.File{stdin, stdout, stderr}}
	p, err := os.StartProcess(path, args, &procattr)
	if err != nil {
		return err
	}

	return p.Release()
}

func tryToConnect(network, address string) (*rpc.Client, error) {
	start := time.Now()
	for {
		client, err := rpc.Dial(network, address)
		if err != nil && time.Since(start) < time.Second {
			continue
		}
		return client, err
	}
}

func cmdAutoComplete(c *rpc.Client) {
	var req AutoCompleteRequest
	req.Filename, req.Data, req.Cursor = prepareFilenameDataCursor()
	req.Context = cache.PackContext(&build.Default)
	req.Source = *g_source
	req.Builtin = *g_builtin
	req.IgnoreCase = *g_ignore_case
	req.UnimportedPackages = *g_unimported_packages
	req.FallbackToSource = *g_fallback_to_source

	var res AutoCompleteReply
	var err error
	if c == nil {
		s := Server{}
		err = s.AutoComplete(&req, &res)
	} else {
		err = c.Call("Server.AutoComplete", &req, &res)
	}
	if err != nil {
		log.Fatal(err)
	}

	fmt := suggest.Formatters[*g_format]
	if fmt == nil {
		fmt = suggest.NiceFormat
	}
	fmt(os.Stdout, res.Candidates, res.Len)
}

func cmdExit(c *rpc.Client) {
	if c == nil {
		return
	}
	var req ExitRequest
	var res ExitReply
	if err := c.Call("Server.Exit", &req, &res); err != nil {
		log.Fatal(err)
	}
}

func prepareFilenameDataCursor() (string, []byte, int) {
	var file []byte
	var err error

	if *g_input != "" {
		file, err = ioutil.ReadFile(*g_input)
	} else {
		file, err = ioutil.ReadAll(os.Stdin)
	}

	if err != nil {
		log.Fatal(err)
	}

	filename := *g_input
	offset := ""
	switch flag.NArg() {
	case 2:
		offset = flag.Arg(1)
	case 3:
		filename = flag.Arg(1) // Override default filename
		offset = flag.Arg(2)
	}

	if filename != "" {
		filename, _ = filepath.Abs(filename)
	}

	cursor := -1
	if offset != "" {
		if offset[0] == 'c' || offset[0] == 'C' {
			cursor, _ = strconv.Atoi(offset[1:])
			cursor = runeToByteOffset(file, cursor)
		} else {
			cursor, _ = strconv.Atoi(offset)
		}
	}

	return filename, file, cursor
}
<commit_msg>Client passes -cache option on to the server it starts<commit_after>package main

import (
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"log"
	"net/rpc"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"runtime/debug"

	"github.com/mdempsky/gocode/internal/cache"
	"github.com/mdempsky/gocode/internal/suggest"
)

func doClient() {
	// Client is a short-lived program.
	// Disable GC to make it faster
	debug.SetGCPercent(-1)

	if *g_debug {
		start := time.Now()
		defer func() {
			elapsed := time.Since(start)
			log.Printf("Elapsed duration: %v\n", elapsed)
		}()
	}

	var command string
	if flag.NArg() > 0 {
		command = flag.Arg(0)
		switch command {
		case "autocomplete", "exit":
			// these are valid commands
		case "close":
			// "close" is an alias for "exit"
			command = "exit"
		default:
			fmt.Printf("gocode: unknown subcommand: %q\nRun 'gocode -help' for usage.\n", command)
			os.Exit(2)
		}
	}

	// client
	var client *rpc.Client
	if *g_sock != "none" {
		addr := *g_addr
		if *g_sock == "unix" {
			addr = getSocketPath()
		}

		var err error
		client, err = rpc.Dial(*g_sock, addr)
		if err != nil {
			if command == "exit" {
				log.Fatal(err)
			}

			if *g_sock == "unix" {
				_ = os.Remove(addr)
			}
			err = tryStartServer()
			if err != nil {
				log.Fatalf("Failed to start server: %s\n", err)
			}
			client, err = tryToConnect(*g_sock, addr)
			if err != nil {
				log.Fatalf("Failed to connect to %q: %s\n", addr, err)
			}
		}
		defer client.Close()
	}

	switch command {
	case "autocomplete":
		cmdAutoComplete(client)
	case "exit":
		cmdExit(client)
	}
}

func tryStartServer() error {
	path := get_executable_filename()
	args := []string{os.Args[0], "-s", "-sock", *g_sock, "-addr", *g_addr}
	if *g_cache {
		args = append(args, "-cache")
	}
	cwd, _ := os.Getwd()

	var err error
	stdin, err := os.Open(os.DevNull)
	if err != nil {
		return err
	}
	stdout, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
	if err != nil {
		return err
	}
	stderr, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
	if err != nil {
		return err
	}

	procattr := os.ProcAttr{Dir: cwd, Env: os.Environ(), Files: []*os.File{stdin, stdout, stderr}}
	p, err := os.StartProcess(path, args, &procattr)
	if err != nil {
		return err
	}

	return p.Release()
}

func tryToConnect(network, address string) (*rpc.Client, error) {
	start := time.Now()
	for {
		client, err := rpc.Dial(network, address)
		if err != nil && time.Since(start) < time.Second {
			continue
		}
		return client, err
	}
}

func cmdAutoComplete(c *rpc.Client) {
	var req AutoCompleteRequest
	req.Filename, req.Data, req.Cursor = prepareFilenameDataCursor()
	req.Context = cache.PackContext(&build.Default)
	req.Source = *g_source
	req.Builtin = *g_builtin
	req.IgnoreCase = *g_ignore_case
	req.UnimportedPackages = *g_unimported_packages
	req.FallbackToSource = *g_fallback_to_source

	var res AutoCompleteReply
	var err error
	if c == nil {
		s := Server{}
		err = s.AutoComplete(&req, &res)
	} else {
		err = c.Call("Server.AutoComplete", &req, &res)
	}
	if err != nil {
		log.Fatal(err)
	}

	fmt := suggest.Formatters[*g_format]
	if fmt == nil {
		fmt = suggest.NiceFormat
	}
	fmt(os.Stdout, res.Candidates, res.Len)
}

func cmdExit(c *rpc.Client) {
	if c == nil {
		return
	}
	var req ExitRequest
	var res ExitReply
	if err := c.Call("Server.Exit", &req, &res); err != nil {
		log.Fatal(err)
	}
}

func prepareFilenameDataCursor() (string, []byte, int) {
	var file []byte
	var err error

	if *g_input != "" {
		file, err = ioutil.ReadFile(*g_input)
	} else {
		file, err = ioutil.ReadAll(os.Stdin)
	}

	if err != nil {
		log.Fatal(err)
	}

	filename := *g_input
	offset := ""
	switch flag.NArg() {
	case 2:
		offset = flag.Arg(1)
	case 3:
		filename = flag.Arg(1) // Override default filename
		offset = flag.Arg(2)
	}

	if filename != "" {
		filename, _ = filepath.Abs(filename)
	}

	cursor := -1
	if offset != "" {
		if offset[0] == 'c' || offset[0] == 'C' {
			cursor, _ = strconv.Atoi(offset[1:])
			cursor = runeToByteOffset(file, cursor)
		} else {
			cursor, _ = strconv.Atoi(offset)
		}
	}

	return filename, file, cursor
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package sse

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"errors"
	"log"
	"net/http"
)

var (
	headerID    = []byte("id:")
	headerData  = []byte("data:")
	headerEvent = []byte("event:")
	headerError = []byte("error:")
)

// Client handles an incoming server stream
type Client struct {
	URL            string
	Connection     *http.Client
	Headers        map[string]string
	EncodingBase64 bool
	EventID        string
}

// NewClient creates a new client
func NewClient(url string) *Client {
	return &Client{
		URL:        url,
		Connection: &http.Client{},
		Headers:    make(map[string]string),
	}
}

// Subscribe to a data stream
func (c *Client) Subscribe(stream string, handler func(msg *Event)) error {
	resp, err := c.request(stream)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	reader := bufio.NewReader(resp.Body)

	for {
		// Read each new line and process the type of event
		line, err := reader.ReadBytes('\n')
		if err != nil {
			return err
		}
		msg := c.processEvent(line)
		if msg != nil {
			handler(msg)
		}
	}
}

// SubscribeChan sends all events to the provided channel
func (c *Client) SubscribeChan(stream string, ch chan *Event) error {
	defer close(ch)

	resp, err := c.request(stream)
	if err != nil {
		return err
	}

	if resp.StatusCode != 200 {
		return errors.New("could not connect to stream")
	}

	defer resp.Body.Close()

	reader := bufio.NewReader(resp.Body)

	for {
		// Read each new line and process the type of event
		line, err := reader.ReadBytes('\n')
		if err != nil {
			close(ch)
			return err
		}
		msg := c.processEvent(line)
		if msg != nil {
			ch <- msg
		}
	}
}

func (c *Client) request(stream string) (*http.Response, error) {
	req, err := http.NewRequest("GET", c.URL, nil)
	if err != nil {
		return nil, err
	}

	// Setup request, specify stream to connect to
	query := req.URL.Query()
	query.Add("stream", stream)
	req.URL.RawQuery = query.Encode()

	req.Header.Set("Cache-Control", "no-cache")
	req.Header.Set("Accept", "text/event-stream")
	req.Header.Set("Connection", "keep-alive")

	if c.EventID != "" {
		req.Header.Set("Last-Event-ID", c.EventID)
	}

	// Add user specified headers
	for k, v := range c.Headers {
		req.Header.Set(k, v)
	}

	return c.Connection.Do(req)
}

func (c *Client) processEvent(msg []byte) *Event {
	var e Event

	switch h := msg; {
	case bytes.Contains(h, headerID):
		e.ID = trimHeader(len(headerID), msg)
	case bytes.Contains(h, headerData):
		e.Data = trimHeader(len(headerData), msg)
	case bytes.Contains(h, headerEvent):
		e.Event = trimHeader(len(headerEvent), msg)
	case bytes.Contains(h, headerError):
		e.Error = trimHeader(len(headerError), msg)
	default:
		return nil
	}

	if len(e.Data) > 0 && c.EncodingBase64 {
		buf := make([]byte, base64.StdEncoding.DecodedLen(len(e.Data)))

		_, err := base64.StdEncoding.Decode(buf, e.Data)
		if err != nil {
			log.Println(err)
		}

		e.Data = buf
	}

	return &e
}

func trimHeader(size int, data []byte) []byte {
	data = data[size:]
	// Remove optional leading whitespace
	if data[0] == 32 {
		data = data[1:]
	}
	// Remove trailing new line
	if data[len(data)-1] == 10 {
		data = data[:len(data)-1]
	}
	return data
}
<commit_msg>SubscribeChan now correctly runs in a goroutine<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package sse

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"errors"
	"log"
	"net/http"
)

var (
	headerID    = []byte("id:")
	headerData  = []byte("data:")
	headerEvent = []byte("event:")
	headerError = []byte("error:")
)

// Client handles an incoming server stream
type Client struct {
	URL            string
	Connection     *http.Client
	Headers        map[string]string
	EncodingBase64 bool
	EventID        string
}

// NewClient creates a new client
func NewClient(url string) *Client {
	return &Client{
		URL:        url,
		Connection: &http.Client{},
		Headers:    make(map[string]string),
	}
}

// Subscribe to a data stream
func (c *Client) Subscribe(stream string, handler func(msg *Event)) error {
	resp, err := c.request(stream)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	reader := bufio.NewReader(resp.Body)

	for {
		// Read each new line and process the type of event
		line, err := reader.ReadBytes('\n')
		if err != nil {
			return err
		}
		msg := c.processEvent(line)
		if msg != nil {
			handler(msg)
		}
	}
}

// SubscribeChan sends all events to the provided channel
func (c *Client) SubscribeChan(stream string, ch chan *Event) error {
	resp, err := c.request(stream)
	if err != nil {
		close(ch)
		return err
	}

	if resp.StatusCode != 200 {
		close(ch)
		return errors.New("could not connect to stream")
	}

	reader := bufio.NewReader(resp.Body)

	go func() {
		for {
			// Read each new line and process the type of event
			line, err := reader.ReadBytes('\n')
			if err != nil {
				resp.Body.Close()
				close(ch)
				return
			}
			msg := c.processEvent(line)
			if msg != nil {
				ch <- msg
			}
		}
	}()

	return nil
}

func (c *Client) request(stream string) (*http.Response, error) {
	req, err := http.NewRequest("GET", c.URL, nil)
	if err != nil {
		return nil, err
	}

	// Setup request, specify stream to connect to
	query := req.URL.Query()
	query.Add("stream", stream)
	req.URL.RawQuery = query.Encode()

	req.Header.Set("Cache-Control", "no-cache")
	req.Header.Set("Accept", "text/event-stream")
	req.Header.Set("Connection", "keep-alive")

	if c.EventID != "" {
		req.Header.Set("Last-Event-ID", c.EventID)
	}

	// Add user specified headers
	for k, v := range c.Headers {
		req.Header.Set(k, v)
	}

	return c.Connection.Do(req)
}

func (c *Client) processEvent(msg []byte) *Event {
	var e Event

	switch h := msg; {
	case bytes.Contains(h, headerID):
		e.ID = trimHeader(len(headerID), msg)
	case bytes.Contains(h, headerData):
		e.Data = trimHeader(len(headerData), msg)
	case bytes.Contains(h, headerEvent):
		e.Event = trimHeader(len(headerEvent), msg)
	case bytes.Contains(h, headerError):
		e.Error = trimHeader(len(headerError), msg)
	default:
		return nil
	}

	if len(e.Data) > 0 && c.EncodingBase64 {
		buf := make([]byte, base64.StdEncoding.DecodedLen(len(e.Data)))

		_, err := base64.StdEncoding.Decode(buf, e.Data)
		if err != nil {
			log.Println(err)
		}

		e.Data = buf
	}

	return &e
}

func trimHeader(size int, data []byte) []byte {
	data = data[size:]
	// Remove optional leading whitespace
	if data[0] == 32 {
		data = data[1:]
	}
	// Remove trailing new line
	if data[len(data)-1] == 10 {
		data = data[:len(data)-1]
	}
	return data
}
<|endoftext|>


<commit_before>// Package itembase gives a thin wrapper around the itembase REST API.
package itembase

import (
	"encoding/json"
)

// Error is a Go representation of the error message sent back by itembase when a
// request results in an error.
type Error struct {
	Message string `json:"message"`
	Code    int    `json:"code"`
}

func (f *Error) Error() string {
	return f.Message
}

// This is the default implementation
type client struct {
	// root is the client's base URL used for all calls.
	root       string
	me         string
	activation string

	// url is the current url to call
	url string

	// auth is authentication token used when making calls.
	// The token is optional and can also be overwritten on an individual
	// call basis via params.
	auth string

	// user is the current shop we're calling for
	user string

	// production environment vs sandbox
	production bool

	// api is the underlying client used to make calls.
	api API

	params  map[string]string
	options Config
}

// New creates a new instance of the default itembase Client implementation.
//
// The options must be non-nil and must provide all OAuth2 credentials and
// configuration for an application registered with the itembase API.
//
// TODO: always use the default API impl, NewClient allows dependency injection
// needed for testing.
func New(options Config, api API) Client {
	if api == nil {
		api = new(itembaseAPI)
	}

	return &client{options: options, production: options.Production, api: api}
}

// NewClient is an alternative Client constructor intended for testing or
// advanced usage, where a custom API implementation can be injected.
func NewClient(root, auth string, options Config, api API) Client {
	if api == nil {
		api = new(itembaseAPI)
	}

	newClient := &client{url: root, root: root, auth: auth, api: api, options: options, production: options.Production}
	newClient.newConf()

	return newClient
}

func (c *client) URL() string {
	return c.url
}

func (c *client) Sandbox() Client {
	c.production = false
	return c
}

func (c *client) User(user string) Client {
	c.auth = c.getUserToken(user).AccessToken
	c.user = user
	c.params = make(map[string]string)
	c.url = c.root + "/users/" + user
	return c
}

func (c *client) GetInto(destination interface{}) error {
	err := c.api.Call("GET", c.url, c.auth, nil, c.params, &destination)
	if err != nil {
		return err
	}
	return nil
}

func (c *client) Get() (destination interface{}, err error) {
	err = c.api.Call("GET", c.url, c.auth, nil, c.params, &destination)
	return
}

func (c *client) GetAllInto(destination interface {
	Add(interface{}) error
	Count() int
	MaxCreatedAt() time.Time
	MaxUpdatedAt() time.Time
}) (err error) {

	var response ItembaseResponse
	DocumentsReceived := 0

	err = c.api.Call("GET", c.url, c.auth, nil, c.params, &response)
	if err != nil {
		return
	}

	for _, document := range response.Documents {
		if destination.Add != nil {
			err = destination.Add(document)
			if err != nil {
				log.Info("Error when adding document", "error", err)
			}
		}
	}

	if response.NumDocumentsFound == response.NumDocumentsReturned {
		return
	} else {
		DocumentsReceived = response.NumDocumentsReturned

		for DocumentsReceived < response.NumDocumentsFound {

			c = c.clientWithNewParam("start_at_document", DocumentsReceived)
			err = c.api.Call("GET", c.url, c.auth, nil, c.params, &response)

			if len(response.Documents) == 0 {
				return
			}

			DocumentsReceived = DocumentsReceived + response.NumDocumentsReturned
			for _, document := range response.Documents {
				if destination.Add != nil {
					destination.Add(document)
				}
			}
		}
	}

	return
}

func (c *client) Found() (count int, err error) {

	var response ItembaseResponse

	d := c.clientWithNewParam("limit", 1)
	err = d.api.Call("GET", d.url, d.auth, nil, d.params, &response)

	if err != nil {
		return
	}

	count = response.NumDocumentsFound
	return
}

func (c *client) Me() (destination User, err error) {
	err = c.api.Call("GET", c.me, c.auth, nil, c.params, &destination)
	return
}

func (c *client) Activate() (destination interface{}, err error) {
	err = c.api.Call("GET", c.activation+"/activate", c.auth, nil, c.params, &destination)
	return
}

func (c *client) Child(path string) Client {
	c.url = c.url + "/" + path
	return c
}

func (c *client) Transactions() Client {
	c.url = c.root + "/users/" + c.user + "/transactions"
	return c
}

func (c *client) Products() Client {
	c.url = c.root + "/users/" + c.user + "/products"
	return c
}

func (c *client) Buyers() Client {
	c.url = c.root + "/users/" + c.user + "/buyers"
	return c
}

func (c *client) Profiles() Client {
	c.url = c.root + "/users/" + c.user + "/profiles"
	return c
}

// These are some shenanigans, golang. Shenanigans I say.
func (c *client) newParamMap(key string, value interface{}) map[string]string {
	ret := make(map[string]string, len(c.params)+1)
	for key, value := range c.params {
		ret[key] = value
	}
	switch value.(type) {
	case string:
		ret[key] = value.(string)
	default:
		jsonVal, _ := json.Marshal(value)
		ret[key] = string(jsonVal)
	}
	return ret
}

func (c *client) clientWithNewParam(key string, value interface{}) *client {
	c.params = c.newParamMap(key, value)
	return c
}

// Query functions.
func (c *client) Select(prop string) Client {
	c.url = c.url + "/" + prop
	return c
}

func (c *client) CreatedAtFrom(value string) Client {
	return c.clientWithNewParam("created_at_from", value)
}

func (c *client) CreatedAtTo(value string) Client {
	return c.clientWithNewParam("created_at_to", value)
}

func (c *client) UpdatedAtFrom(value string) Client {
	return c.clientWithNewParam("updated_at_from", value)
}

func (c *client) UpdatedAtTo(value string) Client {
	return c.clientWithNewParam("updated_at_to", value)
}

func (c *client) Limit(limit uint) Client {
	return c.clientWithNewParam("document_limit", limit)
}

func (c *client) Offset(offset uint) Client {
	return c.clientWithNewParam("start_at_document", offset)
}
<commit_msg>Expect time for time-based methods<commit_after>// Package itembase gives a thin wrapper around the itembase REST API.
package itembase

import (
	"encoding/json"
)

// Error is a Go representation of the error message sent back by itembase when a
// request results in an error.
type Error struct {
	Message string `json:"message"`
	Code    int    `json:"code"`
}

func (f *Error) Error() string {
	return f.Message
}

// This is the default implementation
type client struct {
	// root is the client's base URL used for all calls.
	root       string
	me         string
	activation string

	// url is the current url to call
	url string

	// auth is authentication token used when making calls.
	// The token is optional and can also be overwritten on an individual
	// call basis via params.
	auth string

	// user is the current shop we're calling for
	user string

	// production environment vs sandbox
	production bool

	// api is the underlying client used to make calls.
	api API

	params  map[string]string
	options Config
}

// New creates a new instance of the default itembase Client implementation.
//
// The options must be non-nil and must provide all OAuth2 credentials and
// configuration for an application registered with the itembase API.
//
// TODO: always use the default API impl, NewClient allows dependency injection
// needed for testing.
func New(options Config, api API) Client {
	if api == nil {
		api = new(itembaseAPI)
	}

	return &client{options: options, production: options.Production, api: api}
}

// NewClient is an alternative Client constructor intended for testing or
// advanced usage, where a custom API implementation can be injected.
func NewClient(root, auth string, options Config, api API) Client {
	if api == nil {
		api = new(itembaseAPI)
	}

	newClient := &client{url: root, root: root, auth: auth, api: api, options: options, production: options.Production}
	newClient.newConf()

	return newClient
}

func (c *client) URL() string {
	return c.url
}

func (c *client) Sandbox() Client {
	c.production = false
	return c
}

func (c *client) User(user string) Client {
	c.auth = c.getUserToken(user).AccessToken
	c.user = user
	c.params = make(map[string]string)
	c.url = c.root + "/users/" + user
	return c
}

func (c *client) GetInto(destination interface{}) error {
	err := c.api.Call("GET", c.url, c.auth, nil, c.params, &destination)
	if err != nil {
		return err
	}
	return nil
}

func (c *client) Get() (destination interface{}, err error) {
	err = c.api.Call("GET", c.url, c.auth, nil, c.params, &destination)
	return
}

func (c *client) GetAllInto(destination interface {
	Add(interface{}) error
	Count() int
	MaxCreatedAt() time.Time
	MaxUpdatedAt() time.Time
}) (err error) {

	var response ItembaseResponse
	DocumentsReceived := 0

	err = c.api.Call("GET", c.url, c.auth, nil, c.params, &response)
	if err != nil {
		return
	}

	for _, document := range response.Documents {
		if destination.Add != nil {
			err = destination.Add(document)
			if err != nil {
				log.Info("Error when adding document", "error", err)
			}
		}
	}

	if response.NumDocumentsFound == response.NumDocumentsReturned {
		return
	} else {
		DocumentsReceived = response.NumDocumentsReturned

		for DocumentsReceived < response.NumDocumentsFound {

			c = c.clientWithNewParam("start_at_document", DocumentsReceived)
			err = c.api.Call("GET", c.url, c.auth, nil, c.params, &response)

			if len(response.Documents) == 0 {
				return
			}

			DocumentsReceived = DocumentsReceived + response.NumDocumentsReturned
			for _, document := range response.Documents {
				if destination.Add != nil {
					destination.Add(document)
				}
			}
		}
	}

	return
}

func (c *client) Found() (count int, err error) {

	var response ItembaseResponse

	d := c.clientWithNewParam("limit", 1)
	err = d.api.Call("GET", d.url, d.auth, nil, d.params, &response)

	if err != nil {
		return
	}

	count = response.NumDocumentsFound
	return
}

func (c *client) Me() (destination User, err error) {
	err = c.api.Call("GET", c.me, c.auth, nil, c.params, &destination)
	return
}

func (c *client) Activate() (destination interface{}, err error) {
	err = c.api.Call("GET", c.activation+"/activate", c.auth, nil, c.params, &destination)
	return
}

func (c *client) Child(path string) Client {
	c.url = c.url + "/" + path
	return c
}

func (c *client) Transactions() Client {
	c.url = c.root + "/users/" + c.user + "/transactions"
	return c
}

func (c *client) Products() Client {
	c.url = c.root + "/users/" + c.user + "/products"
	return c
}

func (c *client) Buyers() Client {
	c.url = c.root + "/users/" + c.user + "/buyers"
	return c
}

func (c *client) Profiles() Client {
	c.url = c.root + "/users/" + c.user + "/profiles"
	return c
}

// These are some shenanigans, golang. Shenanigans I say.
func (c *client) newParamMap(key string, value interface{}) map[string]string {
	ret := make(map[string]string, len(c.params)+1)
	for key, value := range c.params {
		ret[key] = value
	}
	switch value.(type) {
	case string:
		ret[key] = value.(string)
	default:
		jsonVal, _ := json.Marshal(value)
		ret[key] = string(jsonVal)
	}
	return ret
}

func (c *client) clientWithNewParam(key string, value interface{}) *client {
	c.params = c.newParamMap(key, value)
	return c
}

// Query functions.
func (c *client) Select(prop string) Client {
	c.url = c.url + "/" + prop
	return c
}

func (c *client) CreatedAtFrom(value time.Time) Client {
	return c.clientWithNewParam("created_at_from", value.Format(time.RFC3339Nano))
}

func (c *client) CreatedAtTo(value time.Time) Client {
	return c.clientWithNewParam("created_at_to", value.Format(time.RFC3339Nano))
}

func (c *client) UpdatedAtFrom(value time.Time) Client {
	return c.clientWithNewParam("updated_at_from", value.Format(time.RFC3339Nano))
}

func (c *client) UpdatedAtTo(value time.Time) Client {
	return c.clientWithNewParam("updated_at_to", value.Format(time.RFC3339Nano))
}

func (c *client) Limit(limit uint) Client {
	return c.clientWithNewParam("document_limit", limit)
}

func (c *client) Offset(offset uint) Client {
	return c.clientWithNewParam("start_at_document", offset)
}
<|endoftext|>


<commit_before>//
// mod.go
// post moderation
//
package srnd

import (
  "bytes"
  "errors"
  "fmt"
  "io"
  "log"
  "net/http"
  "os"
  "strings"
)

// interface for moderation ui
type ModUI interface {

  // channel for daemon to poll for nntp articles from the mod ui
  MessageChan() chan NNTPMessage

  // check if this key is allowed to access
  // return true if it can otherwise false
  CheckKey(privkey string) (bool, error)

  // serve the base page
  ServeModPage(wr http.ResponseWriter, r *http.Request)
  // handle a login POST request
  HandleLogin(wr http.ResponseWriter, r *http.Request)
  // handle a delete article request
  HandleDeletePost(wr http.ResponseWriter, r *http.Request)
  // handle a ban address request
  HandleBanAddress(wr http.ResponseWriter, r *http.Request)
  // handle an unban address request
  HandleUnbanAddress(wr http.ResponseWriter, r *http.Request)
  // handle add a pubkey
  HandleAddPubkey(wr http.ResponseWriter, r *http.Request)
  // handle removing a pubkey
  HandleDelPubkey(wr http.ResponseWriter, r *http.Request)
  // handle key generation
  HandleKeyGen(wr http.ResponseWriter, r *http.Request)
}

type ModEvent interface {
  // turn it into a string for putting into an article
  String() string
  // what type of mod event
  Action() string
  // what reason for the event
  Reason() string
  // what is the event acting on
  Target() string
  // scope of the event, regex of newsgroup
  Scope() string
  // when this mod event expires, unix nano
  Expires() int64
}

type simpleModEvent string

func (self simpleModEvent) String() string {
  return string(self)
}

func (self simpleModEvent) Action() string {
  return strings.Split(string(self), " ")[0]
}

func (self simpleModEvent) Reason() string {
  return ""
}

func (self simpleModEvent) Target() string {
  return strings.Split(string(self), " ")[1]
}

func (self simpleModEvent) Scope() string {
  // TODO: hard coded
  return "overchan.*"
}

func (self simpleModEvent) Expires() int64 {
  // no expiration
  return -1
}

// create an overchan-delete mod event
func overchanDelete(msgid string) ModEvent {
  return simpleModEvent(fmt.Sprintf("delete %s", msgid))
}

// create an overchan-inet-ban mod event
func overchanInetBan(encAddr, key string, expire int64) ModEvent {
  return simpleModEvent(fmt.Sprintf("overchan-inet-ban %s:%s:%d", encAddr, key, expire))
}

// moderation message
// wraps multiple mod events
// is turned into an NNTPMessage later
type ModMessage []ModEvent

// write this mod message's body
func (self ModMessage) WriteTo(wr io.Writer, delim []byte) (err error) {
  // write body
  for _, ev := range(self) {
    _, err = io.WriteString(wr, ev.String())
    _, err = wr.Write(delim)
  }
  return
}

func ParseModEvent(line string) ModEvent {
  return simpleModEvent(line)
}

// wrap mod message in an nntp message
// does not sign
func wrapModMessage(mm ModMessage) NNTPMessage {
  pathname := "nntpchan.censor"
  nntp := nntpArticle{
    headers: make(ArticleHeaders),
  }
  nntp.headers.Set("Newsgroups", "ctl")
  nntp.headers.Set("Content-Type", "text/plain; charset=UTF-8")
  nntp.headers.Set("Message-ID", genMessageID(pathname))
  nntp.headers.Set("Date", timeNowStr())
  nntp.headers.Set("Path", pathname)
  // todo: set these maybe?
  nntp.headers.Set("From", "anon <a@n.on>")
  nntp.headers.Set("Subject", "censor")

  var buff bytes.Buffer
  // crlf delimited
  _ = mm.WriteTo(&buff, []byte{13,10})
  // create plaintext attachment, cut off last 2 bytes
  str := buff.String()
  nntp.message = createPlaintextAttachment(str[:len(str)-2])
  return nntp
}

type ModEngine interface {
  // chan to send the mod engine posts
  // assumes ctl newsgroup only
  MessageChan() chan NNTPMessage

  // delete post of a poster
  DeletePost(msgid string) error
  // ban a cidr
  BanAddress(cidr string) error
  // do we allow this public key to delete?
  AllowDelete(pubkey string) bool
  // do we allow this public key to ban?
  AllowBan(pubkey string) bool
}

type modEngine struct {
  database Database
  store ArticleStore
  chnl chan NNTPMessage
}

func (self modEngine) MessageChan() chan NNTPMessage {
  return self.chnl
}

func (self modEngine) BanAddress(cidr string) (err error) {
  return self.database.BanAddr(cidr)
}

func (self modEngine) DeletePost(msgid string) (err error) {
  hdr := self.store.GetHeaders(msgid)
  if hdr == nil {
    return errors.New("no such message on filesystem: "+msgid)
  } else {
    ref := hdr.Get("References", "")
    var delposts []string
    if ref == "" {
      // is root post
      // delete replies too
      repls := self.database.GetThreadReplies(msgid, 0)
      if repls == nil {
        log.Println("cannot get thread replies for", msgid)
      } else {
        delposts = append(delposts, repls...)
      }
      // delete thread presence
      self.database.DeleteThread(msgid)
    }
    delposts = append(delposts, msgid)
    // get list of files to delete
    var delfiles []string
    for _, delmsg := range delposts {
      article := self.store.GetFilename(delmsg)
      delfiles = append(delfiles, article)
      // get attachments for post
      atts := self.database.GetPostAttachments(delmsg)
      if atts != nil {
        for _, att := range(atts) {
          img := self.store.AttachmentFilepath(att)
          thm := self.store.ThumbnailFilepath(att)
          delfiles = append(delfiles, img, thm)
        }
      }
      // delete article from database
      self.database.DeleteArticle(delmsg)
    }
    // delete all files
    for _, f := range delfiles {
      log.Printf("delete file: %s", f)
      os.Remove(f)
    }
  }
  return nil
}

// TODO: permissions
func (self modEngine) AllowBan(pubkey string) bool {
  return self.database.CheckModPubkeyGlobal(pubkey)
}
// TODO: permissions
func (self modEngine) AllowDelete(pubkey string) bool {
  return self.database.CheckModPubkeyGlobal(pubkey)
}

// run a mod engine logic mainloop
func RunModEngine(mod ModEngine) {

  chnl := mod.MessageChan()
  for {
    nntp := <- chnl
    // sanity check
    if nntp.Newsgroup() == "ctl" {
      inner_nntp := nntp.Signed()
      if inner_nntp != nil {
        // okay this message should be good
        pubkey := nntp.Pubkey()
        for _, line := range strings.Split(inner_nntp.Message(), "\n") {
          ev := ParseModEvent(line)
          action := ev.Action()
          if action == "delete" {
            msgid := ev.Target()
            // this is a delete action
            if mod.AllowDelete(pubkey) {
              mod.DeletePost(msgid)
            } else {
              log.Printf("pubkey=%s will not delete %s not trusted", pubkey, msgid)
            }
          } else if action == "overchan-inet-ban" {
            // ban action
            target := ev.Target()
            parts := strings.Split(target, ":")
            if len(parts) == 3 {
              encaddr, key := parts[0], parts[1]
              cidr := decAddr(encaddr, key)
              if cidr == "" {
                log.Println("failed to decrypt inet ban")
              } else if mod.AllowBan(pubkey) {
                mod.BanAddress(cidr)
              }
            } else {
              log.Printf("invalid overchan-inet-ban: target=%s", target)
            }
          }
        }
      }
    }
  }
}
<commit_msg>mod system remote delete works<commit_after>//
// mod.go
// post moderation
//
package srnd

import (
  "bytes"
  "errors"
  "fmt"
  "io"
  "log"
  "net/http"
  "os"
  "strings"
)

// interface for moderation ui
type ModUI interface {

  // channel for daemon to poll for nntp articles from the mod ui
  MessageChan() chan NNTPMessage

  // check if this key is allowed to access
  // return true if it can otherwise false
  CheckKey(privkey string) (bool, error)

  // serve the base page
  ServeModPage(wr http.ResponseWriter, r *http.Request)
  // handle a login POST request
  HandleLogin(wr http.ResponseWriter, r *http.Request)
  // handle a delete article request
  HandleDeletePost(wr http.ResponseWriter, r *http.Request)
  // handle a ban address request
  HandleBanAddress(wr http.ResponseWriter, r *http.Request)
  // handle an unban address request
  HandleUnbanAddress(wr http.ResponseWriter, r *http.Request)
  // handle add a pubkey
  HandleAddPubkey(wr http.ResponseWriter, r *http.Request)
  // handle removing a pubkey
  HandleDelPubkey(wr http.ResponseWriter, r *http.Request)
  // handle key generation
  HandleKeyGen(wr http.ResponseWriter, r *http.Request)
}

type ModEvent interface {
  // turn it into a string for putting into an article
  String() string
  // what type of mod event
  Action() string
  // what reason for the event
  Reason() string
  // what is the event acting on
  Target() string
  // scope of the event, regex of newsgroup
  Scope() string
  // when this mod event expires, unix nano
  Expires() int64
}

type simpleModEvent string

func (self simpleModEvent) String() string {
  return string(self)
}

func (self simpleModEvent) Action() string {
  return strings.Split(string(self), " ")[0]
}

func (self simpleModEvent) Reason() string {
  return ""
}

func (self simpleModEvent) Target() string {
  return strings.Split(string(self), " ")[1]
}

func (self simpleModEvent) Scope() string {
  // TODO: hard coded
  return "overchan.*"
}

func (self simpleModEvent) Expires() int64 {
  // no expiration
  return -1
}

// create an overchan-delete mod event
func overchanDelete(msgid string) ModEvent {
  return simpleModEvent(fmt.Sprintf("delete %s", msgid))
}

// create an overchan-inet-ban mod event
func overchanInetBan(encAddr, key string, expire int64) ModEvent {
  return simpleModEvent(fmt.Sprintf("overchan-inet-ban %s:%s:%d", encAddr, key, expire))
}

// moderation message
// wraps multiple mod events
// is turned into an NNTPMessage later
type ModMessage []ModEvent

// write this mod message's body
func (self ModMessage) WriteTo(wr io.Writer, delim []byte) (err error) {
  // write body
  for _, ev := range(self) {
    _, err = io.WriteString(wr, ev.String())
    _, err = wr.Write(delim)
  }
  return
}

func ParseModEvent(line string) ModEvent {
  return simpleModEvent(line)
}

// wrap mod message in an nntp message
// does not sign
func wrapModMessage(mm ModMessage) NNTPMessage {
  pathname := "nntpchan.censor"
  nntp := nntpArticle{
    headers: make(ArticleHeaders),
  }
  nntp.headers.Set("Newsgroups", "ctl")
  nntp.headers.Set("Content-Type", "text/plain; charset=UTF-8")
  nntp.headers.Set("Message-ID", genMessageID(pathname))
  nntp.headers.Set("Date", timeNowStr())
  nntp.headers.Set("Path", pathname)
  // todo: set these maybe?
  nntp.headers.Set("From", "anon <a@n.on>")
  nntp.headers.Set("Subject", "censor")

  var buff bytes.Buffer
  // crlf delimited
  _ = mm.WriteTo(&buff, []byte{13,10})
  // create plaintext attachment, cut off last 2 bytes
  str := buff.String()
  nntp.message = createPlaintextAttachment(str[:len(str)-2])
  return nntp
}

type ModEngine interface {
  // chan to send the mod engine posts
  // assumes ctl newsgroup only
  MessageChan() chan NNTPMessage

  // delete post of a poster
  DeletePost(msgid string) error
  // ban a cidr
  BanAddress(cidr string) error
  // do we allow this public key to delete?
  AllowDelete(pubkey string) bool
  // do we allow this public key to ban?
  AllowBan(pubkey string) bool
}

type modEngine struct {
  database Database
  store ArticleStore
  chnl chan NNTPMessage
}

func (self modEngine) MessageChan() chan NNTPMessage {
  return self.chnl
}

func (self modEngine) BanAddress(cidr string) (err error) {
  return self.database.BanAddr(cidr)
}

func (self modEngine) DeletePost(msgid string) (err error) {
  hdr := self.store.GetHeaders(msgid)
  if hdr == nil {
    return errors.New("no such message on filesystem: "+msgid)
  } else {
    ref := hdr.Get("References", "")
    var delposts []string
    if ref == "" {
      // is root post
      // delete replies too
      repls := self.database.GetThreadReplies(msgid, 0)
      if repls == nil {
        log.Println("cannot get thread replies for", msgid)
      } else {
        delposts = append(delposts, repls...)
      }
      // delete thread presence
      self.database.DeleteThread(msgid)
    }
    delposts = append(delposts, msgid)
    // get list of files to delete
    var delfiles []string
    for _, delmsg := range delposts {
      article := self.store.GetFilename(delmsg)
      delfiles = append(delfiles, article)
      // get attachments for post
      atts := self.database.GetPostAttachments(delmsg)
      if atts != nil {
        for _, att := range(atts) {
          img := self.store.AttachmentFilepath(att)
          thm := self.store.ThumbnailFilepath(att)
          delfiles = append(delfiles, img, thm)
        }
      }
      // delete article from database
      self.database.DeleteArticle(delmsg)
    }
    // delete all files
    for _, f := range delfiles {
      log.Printf("delete file: %s", f)
      os.Remove(f)
    }
  }
  return nil
}

// TODO: permissions
func (self modEngine) AllowBan(pubkey string) bool {
  return self.database.CheckModPubkeyGlobal(pubkey)
}
// TODO: permissions
func (self modEngine) AllowDelete(pubkey string) bool {
  return self.database.CheckModPubkeyGlobal(pubkey)
}

// run a mod engine logic mainloop
func RunModEngine(mod ModEngine) {

  chnl := mod.MessageChan()
  for {
    nntp := <- chnl
    // sanity check
    if nntp.Newsgroup() == "ctl" {
      inner_nntp := nntp.Signed()
      if inner_nntp != nil {
        // okay this message should be good
        pubkey := nntp.Pubkey()
        for _, line := range strings.Split(inner_nntp.Message(), "\n") {
          line = strings.Trim(line, "\r\t\n")
          ev := ParseModEvent(line)
          action := ev.Action()
          if action == "delete" {
            msgid := ev.Target()
            // this is a delete action
            if mod.AllowDelete(pubkey) {
              err := mod.DeletePost(msgid)
              if err != nil {
                log.Println(msgid, err)
              }
            } else {
              log.Printf("pubkey=%s will not delete %s not trusted", pubkey, msgid)
            }
          } else if action == "overchan-inet-ban" {
            // ban action
            target := ev.Target()
            parts := strings.Split(target, ":")
            if len(parts) == 3 {
              encaddr, key := parts[0], parts[1]
              cidr := decAddr(encaddr, key)
              if cidr == "" {
                log.Println("failed to decrypt inet ban")
              } else if mod.AllowBan(pubkey) {
                err := mod.BanAddress(cidr)
                if err != nil {
                  log.Println(cidr, err)
                }
              }
            } else {
              log.Printf("invalid overchan-inet-ban: target=%s", target)
            }
          }
        }
      }
    }
  }
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strings"

	"github.com/asciimoo/colly"
)

func main() {
	c := colly.NewCollector()

	c.AllowedDomains = []string{"coursera.org", "www.coursera.org"}
	courses := make(map[string]string)

	c.OnHTML("a[href]", func(e *colly.HTMLElement) {
		if e.Attr("class") == "Button_1qxkboh-o_O-primary_cv02ee-o_O-md_28awn8-o_O-primaryLink_109aggg" {
			return
		}
		link := e.Attr("href")
		if !strings.HasPrefix(link, "/browse") || strings.Index(link, "=signup") > -1 || strings.Index(link, "=login") > -1 {
			return
		}
		if strings.Count("link", "/") > 2 {
			return
		}
		e.Request.Visit(link)
	})

	c.OnRequest(func(r *colly.Request) {
		log.Println("visiting", r.URL.String())
	})

	c.OnHTML(`a[name]`, func(e *colly.HTMLElement) {
		courses[e.Request.AbsoluteURL(e.Attr("href"))] = e.Text
	})

	c.Visit("https://coursera.org/browse")

	for url, title := range courses {
		fmt.Println(url, "-", title)
	}
}
<commit_msg>[fix] remove buggy coursera url filter<commit_after>package main

import (
	"fmt"
	"log"
	"strings"

	"github.com/asciimoo/colly"
)

func main() {
	c := colly.NewCollector()

	c.AllowedDomains = []string{"coursera.org", "www.coursera.org"}
	courses := make(map[string]string)

	c.OnHTML("a[href]", func(e *colly.HTMLElement) {
		if e.Attr("class") == "Button_1qxkboh-o_O-primary_cv02ee-o_O-md_28awn8-o_O-primaryLink_109aggg" {
			return
		}
		link := e.Attr("href")
		if !strings.HasPrefix(link, "/browse") || strings.Index(link, "=signup") > -1 || strings.Index(link, "=login") > -1 {
			return
		}
		e.Request.Visit(link)
	})

	c.OnRequest(func(r *colly.Request) {
		log.Println("visiting", r.URL.String())
	})

	c.OnHTML(`a[name]`, func(e *colly.HTMLElement) {
		courses[e.Request.AbsoluteURL(e.Attr("href"))] = e.Text
	})

	c.Visit("https://coursera.org/browse")

	for url, title := range courses {
		fmt.Println(url, "-", title)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/hprose/hprose-golang/rpc"
)

type event struct{}

func (event) OnSubscribe(topic string, id string, service rpc.Service) {
	fmt.Println("client " + id + " subscribe topic: " + topic)
}

func (event) OnUnsubscribe(topic string, id string, service rpc.Service) {
	fmt.Println("client " + id + " unsubscribe topic: " + topic)
}

func main() {
	server := rpc.NewTCPServer("tcp4://0.0.0.0:2016/")
	server.Publish("time", 0, 0)
	server.Event = event{}
	var timer *time.Timer
	timer = time.AfterFunc(1*1000*1000*1000, func() {
		server.Broadcast("time", time.Now().String(), func(sended []string) {
			fmt.Println(sended)
		})
		timer.Reset(1 * 1000 * 1000 * 1000)
	})
	server.Start()
}
<commit_msg>Update time_push_server.go<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/hprose/hprose-golang/rpc"
)

type event struct{}

func (event) OnSubscribe(topic string, id string, service rpc.Service) {
	fmt.Println("client " + id + " subscribe topic: " + topic)
}

func (event) OnUnsubscribe(topic string, id string, service rpc.Service) {
	fmt.Println("client " + id + " unsubscribe topic: " + topic)
}

func main() {
	server := rpc.NewTCPServer("tcp4://0.0.0.0:2016/")
	server.Publish("time", 0, 0)
	server.Event = event{}
	var timer *time.Timer
	timer = time.AfterFunc(1*time.Second, func() {
		server.Broadcast("time", time.Now().String(), func(sended []string) {
			fmt.Println(sended)
		})
		timer.Reset(1 * time.Second)
	})
	server.Start()
}
<|endoftext|>


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

package captcha

import (
	"testing"

	"github.com/astaxie/beego/utils"
)

type byteCounter struct {
	n int64
}

func (bc *byteCounter) Write(b []byte) (int, error) {
	bc.n += int64(len(b))
	return len(b), nil
}

func BenchmarkNewImage(b *testing.B) {
	b.StopTimer()
	d := utils.RandomCreateBytes(challengeNums, defaultChars...)
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		NewImage(d, stdWidth, stdHeight)
	}
}

func BenchmarkImageWriteTo(b *testing.B) {
	b.StopTimer()
	d := utils.RandomCreateBytes(challengeNums, defaultChars...)
	b.StartTimer()
	counter := &byteCounter{}
	for i := 0; i < b.N; i++ {
		img := NewImage(d, stdWidth, stdHeight)
		img.WriteTo(counter)
		b.SetBytes(counter.n)
		counter.n = 0
	}
}
<commit_msg>Fix test case<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package captcha

import (
	"testing"

	"github.com/gogits/gogs/modules/base"
)

type byteCounter struct {
	n int64
}

func (bc *byteCounter) Write(b []byte) (int, error) {
	bc.n += int64(len(b))
	return len(b), nil
}

func BenchmarkNewImage(b *testing.B) {
	b.StopTimer()
	d := base.RandomCreateBytes(challengeNums, defaultChars...)
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		NewImage(d, stdWidth, stdHeight)
	}
}

func BenchmarkImageWriteTo(b *testing.B) {
	b.StopTimer()
	d := base.RandomCreateBytes(challengeNums, defaultChars...)
	b.StartTimer()
	counter := &byteCounter{}
	for i := 0; i < b.N; i++ {
		img := NewImage(d, stdWidth, stdHeight)
		img.WriteTo(counter)
		b.SetBytes(counter.n)
		counter.n = 0
	}
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 Zack Scholl. All rights reserved.
// Use of this source code is governed by a AGPL
// license that can be found in the LICENSE file.

// routes.go contains the functions that handle the web page views

package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"

	"github.com/gin-gonic/contrib/sessions"
	"github.com/gin-gonic/gin"
)

// slash returns the dashboard, if logged in, else it redirects to login page.
func slash(c *gin.Context) {
	var group string
	loginGroup := sessions.Default(c)
	groupCookie := loginGroup.Get("group")
	if groupCookie == nil {
		c.Redirect(302, "/login")
	} else {
		group = groupCookie.(string)
		c.Redirect(302, "/dashboard/"+group)
	}
}

// slashLogin handles a login with url parameter and returns dashboard if successful, else login.
func slashLogin(c *gin.Context) {
	loginGroup := sessions.Default(c)
	group := c.DefaultQuery("group", "noneasdf")
	if group == "noneasdf" {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{})
	} else {
		loginGroup.Set("group", group)
		loginGroup.Save()
		c.Redirect(302, "/dashboard/"+group)
	}
}

// slashLogin handles a POST login and returns dashboard if successful, else login.
func slashLoginPOST(c *gin.Context) {
	loginGroup := sessions.Default(c)
	group := strings.ToLower(c.PostForm("group"))
	if _, err := os.Stat(path.Join("data", group+".db")); err == nil {
		loginGroup.Set("group", group)
		loginGroup.Save()
		c.Redirect(302, "/dashboard/"+group)
	} else {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "Incorrect login.",
		})
	}
}

// slashLogout handles a logout
func slashLogout(c *gin.Context) {
	var group string
	loginGroup := sessions.Default(c)
	groupCookie := loginGroup.Get("group")
	if groupCookie == nil {
		c.Redirect(302, "/login")
	} else {
		group = groupCookie.(string)
		fmt.Println(group)
		loginGroup.Clear()
		loginGroup.Save()
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"Message": "You are now logged out.",
		})
	}
}

// slashDashboard displays the dashboard
func slashDashboard(c *gin.Context) {
	// skipUsers := c.DefaultQuery("skip", "")
	// skipAllUsers := false
	// if len(skipUsers) > 0 {
	// 	skipAllUsers = true
	// }
	filterUser := c.DefaultQuery("user", "")
	filterUsers := c.DefaultQuery("users", "")
	filterUserMap := make(map[string]bool)
	if len(filterUser) > 0 {
		u := strings.Replace(strings.TrimSpace(filterUser), ":", "", -1)
		filterUserMap[u] = true
	}
	if len(filterUsers) > 0 {
		for _, user := range strings.Split(filterUsers, ",") {
			u := strings.Replace(strings.TrimSpace(user), ":", "", -1)
			filterUserMap[u] = true
		}
	}
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "First download the app or CLI program to insert some fingerprints.",
		})
		return
	}
	ps, _ := openParameters(group)
	users := getUsers(group)
	people := make(map[string]UserPositionJSON)
	for i, user := range users {
		_, ok := filterUserMap[user]
		if len(filterUserMap) == 0 || ok {
			if i > 10 {
				people[user] = UserPositionJSON{}
			} else {
				people[user] = getCurrentPositionOfUser(group, user)
			}
		}
	}
	type DashboardData struct {
		Networks         []string
		Locations        map[string][]string
		LocationAccuracy map[string]int
		LocationCount    map[string]int
		Mixin            map[string]float64
		VarabilityCutoff map[string]float64
		Users            map[string]UserPositionJSON
	}
	var dash DashboardData
	dash.Networks = []string{}
	dash.Locations = make(map[string][]string)
	dash.LocationAccuracy = make(map[string]int)
	dash.LocationCount = make(map[string]int)
	dash.Mixin = make(map[string]float64)
	dash.VarabilityCutoff = make(map[string]float64)
	for n := range ps.NetworkLocs {
		dash.Mixin[n] = ps.Priors[n].Special["MixIn"]
		dash.VarabilityCutoff[n] = ps.Priors[n].Special["VarabilityCutoff"]
		dash.Networks = append(dash.Networks, n)
		dash.Locations[n] = []string{}
		for loc := range ps.NetworkLocs[n] {
			dash.Locations[n] = append(dash.Locations[n], loc)
			dash.LocationAccuracy[loc] = ps.Results[n].Accuracy[loc]
			dash.LocationCount[loc] = ps.Results[n].TotalLocations[loc]
		}
	}
	c.HTML(http.StatusOK, "dashboard.tmpl", gin.H{
		"Message": RuntimeArgs.Message,
		"Group":   group,
		"Dash":    dash,
		"Users":   people,
	})
}

// slash Location returns location (to be deprecated)
func slashLocation(c *gin.Context) {
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.JSON(http.StatusOK, gin.H{"success": "false", "message": "First download the app or CLI program to insert some fingerprints."})
		return
	}
	user := c.Param("user")
	userJSON := getCurrentPositionOfUser(group, user)
	c.JSON(http.StatusOK, userJSON)
}

// slashExplore returns a chart of the data
func slashExplore(c *gin.Context) {
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "First download the app or CLI program to insert some fingerprints.",
		})
		return
	}
	network := c.Param("network")
	location := c.Param("location")
	ps, _ := openParameters(group)
	// TODO: check if network and location exists in the ps, if not return 404
	datas := []template.JS{}
	names := []template.JS{}
	indexNames := []template.JS{}
	// Sort locations
	macs := []string{}
	for m := range ps.Priors[network].P[location] {
		if float64(ps.MacVariability[m]) > ps.Priors[network].Special["VarabilityCutoff"] {
			macs = append(macs, m)
		}
	}
	sort.Strings(macs)
	it := 0
	for _, m := range macs {
		n := ps.Priors[network].P[location][m]
		names = append(names, template.JS(string(m)))
		jsonByte, _ := json.Marshal(n)
		datas = append(datas, template.JS(string(jsonByte)))
		indexNames = append(indexNames, template.JS(strconv.Itoa(it)))
		it++
	}
	rsiRange, _ := json.Marshal(RssiRange)
	c.HTML(http.StatusOK, "plot.tmpl", gin.H{
		"RssiRange":  template.JS(string(rsiRange)),
		"Datas":      datas,
		"Names":      names,
		"IndexNames": indexNames,
	})
}

// slashExplore returns a chart of the data (canvas.js)
func slashExplore2(c *gin.Context) {
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "First download the app or CLI program to insert some fingerprints.",
		})
		return
	}

	network := c.Param("network")
	location := c.Param("location")
	ps, _ := openParameters(group)
	fmt.Println(ps.UniqueLocs)
	lookUpLocation := false

	for _, loc := range ps.UniqueLocs {
		if location == loc {
			lookUpLocation = true
		}
	}

	type macDatum struct {
		Name   string    `json:"name"`
		Points []float32 `json:"data"`
	}

	type macDatas struct {
		Macs []macDatum `json:"macs"`
	}

	var data macDatas
	data.Macs = []macDatum{}

	if lookUpLocation {
		// Sort locations
		macs := []string{}
		for m := range ps.Priors[network].P[location] {
			if float64(ps.MacVariability[m]) > ps.Priors[network].Special["VarabilityCutoff"] {
				macs = append(macs, m)
			}
		}
		sort.Strings(macs)

		for _, m := range macs {
			n := ps.Priors[network].P[location][m]
			data.Macs = append(data.Macs, macDatum{Name: m, Points: n})
		}
	} else {
		m := location
		for loc := range ps.Priors[network].P {
			Debug.Println(loc, m)
			n := ps.Priors[network].P[loc][m]
			data.Macs = append(data.Macs, macDatum{Name: strings.Replace(loc, " ", "%20", -1), Points: n})
		}
	}

	c.HTML(http.StatusOK, "plot2.tmpl", gin.H{
		"Data":    data,
		"Rssi":    RssiRange,
		"Title":   group + "/" + network + "/" + location,
		"Group":   strings.Replace(group, " ", "%20", -1),
		"Network": strings.Replace(network, " ", "%20", -1),
		"Legend":  !lookUpLocation,
	})
}

// slashPie returns a Pie chart
func slashPie(c *gin.Context) {
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "First download the app or CLI program to insert some fingerprints.",
		})
		return
	}

	network := c.Param("network")
	location := c.Param("location")
	ps, _ := openParameters(group)
	vals := []int{}
	names := []string{}
	fmt.Println(ps.Results[network].Guess[location])
	for guessloc, val := range ps.Results[network].Guess[location] {
		names = append(names, guessloc)
		vals = append(vals, val)
	}
	namesJSON, _ := json.Marshal(names)
	valsJSON, _ := json.Marshal(vals)
	c.HTML(http.StatusOK, "pie.tmpl", gin.H{
		"Names": template.JS(namesJSON),
		"Vals":  template.JS(valsJSON),
	})
}
<commit_msg>Limit the number of users<commit_after>// Copyright 2015-2016 Zack Scholl. All rights reserved.
// Use of this source code is governed by a AGPL
// license that can be found in the LICENSE file.

// routes.go contains the functions that handle the web page views

package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"

	"github.com/gin-gonic/contrib/sessions"
	"github.com/gin-gonic/gin"
)

// slash returns the dashboard, if logged in, else it redirects to login page.
func slash(c *gin.Context) {
	var group string
	loginGroup := sessions.Default(c)
	groupCookie := loginGroup.Get("group")
	if groupCookie == nil {
		c.Redirect(302, "/login")
	} else {
		group = groupCookie.(string)
		c.Redirect(302, "/dashboard/"+group)
	}
}

// slashLogin handles a login with url parameter and returns dashboard if successful, else login.
func slashLogin(c *gin.Context) {
	loginGroup := sessions.Default(c)
	group := c.DefaultQuery("group", "noneasdf")
	if group == "noneasdf" {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{})
	} else {
		loginGroup.Set("group", group)
		loginGroup.Save()
		c.Redirect(302, "/dashboard/"+group)
	}
}

// slashLogin handles a POST login and returns dashboard if successful, else login.
func slashLoginPOST(c *gin.Context) {
	loginGroup := sessions.Default(c)
	group := strings.ToLower(c.PostForm("group"))
	if _, err := os.Stat(path.Join("data", group+".db")); err == nil {
		loginGroup.Set("group", group)
		loginGroup.Save()
		c.Redirect(302, "/dashboard/"+group)
	} else {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "Incorrect login.",
		})
	}
}

// slashLogout handles a logout
func slashLogout(c *gin.Context) {
	var group string
	loginGroup := sessions.Default(c)
	groupCookie := loginGroup.Get("group")
	if groupCookie == nil {
		c.Redirect(302, "/login")
	} else {
		group = groupCookie.(string)
		fmt.Println(group)
		loginGroup.Clear()
		loginGroup.Save()
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"Message": "You are now logged out.",
		})
	}
}

// slashDashboard displays the dashboard
func slashDashboard(c *gin.Context) {
	// skipUsers := c.DefaultQuery("skip", "")
	// skipAllUsers := false
	// if len(skipUsers) > 0 {
	// 	skipAllUsers = true
	// }
	filterUser := c.DefaultQuery("user", "")
	filterUsers := c.DefaultQuery("users", "")
	filterUserMap := make(map[string]bool)
	if len(filterUser) > 0 {
		u := strings.Replace(strings.TrimSpace(filterUser), ":", "", -1)
		filterUserMap[u] = true
	}
	if len(filterUsers) > 0 {
		for _, user := range strings.Split(filterUsers, ",") {
			u := strings.Replace(strings.TrimSpace(user), ":", "", -1)
			filterUserMap[u] = true
		}
	}
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "First download the app or CLI program to insert some fingerprints.",
		})
		return
	}
	ps, _ := openParameters(group)
	users := getUsers(group)
	people := make(map[string]UserPositionJSON)
	for i, user := range users {
		_, ok := filterUserMap[user]
		if len(filterUserMap) == 0 || ok {
			people[user] = getCurrentPositionOfUser(group, user)
		}
	}
	type DashboardData struct {
		Networks         []string
		Locations        map[string][]string
		LocationAccuracy map[string]int
		LocationCount    map[string]int
		Mixin            map[string]float64
		VarabilityCutoff map[string]float64
		Users            map[string]UserPositionJSON
	}
	var dash DashboardData
	dash.Networks = []string{}
	dash.Locations = make(map[string][]string)
	dash.LocationAccuracy = make(map[string]int)
	dash.LocationCount = make(map[string]int)
	dash.Mixin = make(map[string]float64)
	dash.VarabilityCutoff = make(map[string]float64)
	for n := range ps.NetworkLocs {
		dash.Mixin[n] = ps.Priors[n].Special["MixIn"]
		dash.VarabilityCutoff[n] = ps.Priors[n].Special["VarabilityCutoff"]
		dash.Networks = append(dash.Networks, n)
		dash.Locations[n] = []string{}
		for loc := range ps.NetworkLocs[n] {
			dash.Locations[n] = append(dash.Locations[n], loc)
			dash.LocationAccuracy[loc] = ps.Results[n].Accuracy[loc]
			dash.LocationCount[loc] = ps.Results[n].TotalLocations[loc]
		}
	}
	c.HTML(http.StatusOK, "dashboard.tmpl", gin.H{
		"Message": RuntimeArgs.Message,
		"Group":   group,
		"Dash":    dash,
		"Users":   people,
	})
}

// slash Location returns location (to be deprecated)
func slashLocation(c *gin.Context) {
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.JSON(http.StatusOK, gin.H{"success": "false", "message": "First download the app or CLI program to insert some fingerprints."})
		return
	}
	user := c.Param("user")
	userJSON := getCurrentPositionOfUser(group, user)
	c.JSON(http.StatusOK, userJSON)
}

// slashExplore returns a chart of the data
func slashExplore(c *gin.Context) {
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "First download the app or CLI program to insert some fingerprints.",
		})
		return
	}
	network := c.Param("network")
	location := c.Param("location")
	ps, _ := openParameters(group)
	// TODO: check if network and location exists in the ps, if not return 404
	datas := []template.JS{}
	names := []template.JS{}
	indexNames := []template.JS{}
	// Sort locations
	macs := []string{}
	for m := range ps.Priors[network].P[location] {
		if float64(ps.MacVariability[m]) > ps.Priors[network].Special["VarabilityCutoff"] {
			macs = append(macs, m)
		}
	}
	sort.Strings(macs)
	it := 0
	for _, m := range macs {
		n := ps.Priors[network].P[location][m]
		names = append(names, template.JS(string(m)))
		jsonByte, _ := json.Marshal(n)
		datas = append(datas, template.JS(string(jsonByte)))
		indexNames = append(indexNames, template.JS(strconv.Itoa(it)))
		it++
	}
	rsiRange, _ := json.Marshal(RssiRange)
	c.HTML(http.StatusOK, "plot.tmpl", gin.H{
		"RssiRange":  template.JS(string(rsiRange)),
		"Datas":      datas,
		"Names":      names,
		"IndexNames": indexNames,
	})
}

// slashExplore returns a chart of the data (canvas.js)
func slashExplore2(c *gin.Context) {
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "First download the app or CLI program to insert some fingerprints.",
		})
		return
	}

	network := c.Param("network")
	location := c.Param("location")
	ps, _ := openParameters(group)
	fmt.Println(ps.UniqueLocs)
	lookUpLocation := false

	for _, loc := range ps.UniqueLocs {
		if location == loc {
			lookUpLocation = true
		}
	}

	type macDatum struct {
		Name   string    `json:"name"`
		Points []float32 `json:"data"`
	}

	type macDatas struct {
		Macs []macDatum `json:"macs"`
	}

	var data macDatas
	data.Macs = []macDatum{}

	if lookUpLocation {
		// Sort locations
		macs := []string{}
		for m := range ps.Priors[network].P[location] {
			if float64(ps.MacVariability[m]) > ps.Priors[network].Special["VarabilityCutoff"] {
				macs = append(macs, m)
			}
		}
		sort.Strings(macs)

		for _, m := range macs {
			n := ps.Priors[network].P[location][m]
			data.Macs = append(data.Macs, macDatum{Name: m, Points: n})
		}
	} else {
		m := location
		for loc := range ps.Priors[network].P {
			Debug.Println(loc, m)
			n := ps.Priors[network].P[loc][m]
			data.Macs = append(data.Macs, macDatum{Name: strings.Replace(loc, " ", "%20", -1), Points: n})
		}
	}

	c.HTML(http.StatusOK, "plot2.tmpl", gin.H{
		"Data":    data,
		"Rssi":    RssiRange,
		"Title":   group + "/" + network + "/" + location,
		"Group":   strings.Replace(group, " ", "%20", -1),
		"Network": strings.Replace(network, " ", "%20", -1),
		"Legend":  !lookUpLocation,
	})
}

// slashPie returns a Pie chart
func slashPie(c *gin.Context) {
	group := c.Param("group")
	if _, err := os.Stat(path.Join(RuntimeArgs.SourcePath, group+".db")); os.IsNotExist(err) {
		c.HTML(http.StatusOK, "login.tmpl", gin.H{
			"ErrorMessage": "First download the app or CLI program to insert some fingerprints.",
		})
		return
	}

	network := c.Param("network")
	location := c.Param("location")
	ps, _ := openParameters(group)
	vals := []int{}
	names := []string{}
	fmt.Println(ps.Results[network].Guess[location])
	for guessloc, val := range ps.Results[network].Guess[location] {
		names = append(names, guessloc)
		vals = append(vals, val)
	}
	namesJSON, _ := json.Marshal(names)
	valsJSON, _ := json.Marshal(vals)
	c.HTML(http.StatusOK, "pie.tmpl", gin.H{
		"Names": template.JS(namesJSON),
		"Vals":  template.JS(valsJSON),
	})
}
<|endoftext|>


<commit_before><commit_msg>Capture job errors in Raven/Sentry<commit_after><|endoftext|>


<commit_before><commit_msg>atg. light cleanup<commit_after><|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"net/http"
	"github.com/spf13/cobra"
  "github.com/libgit2/git2go"

)

var httpClient *http.Client

var meCmd = &cobra.Command{
	Use:   "me",
	Short: "A brief description of your command",
	Long: `to quickly create a Cobra application.`,

	Run: func(cmd *cobra.Command, args []string) {

		apiUrl:= "/api/v2/users/myself"
		endpoint:= Endpoint(apiUrl)

		fmt.Println(endpoint)
		fmt.Println(git2go.Name())
	},
}

func init() {
	RootCmd.AddCommand(meCmd)
}
<commit_msg>Working api<commit_after>package cmd

import (
"fmt"
"net/http"
"github.com/spf13/cobra"
// "github.com/libgit2/git2go"

"io/ioutil"
"log"
"os"
)

var httpClient *http.Client

var meCmd = &cobra.Command{
	Use:   "me",
	Short: "A brief description of your command",
	Long: `to quickly create a Cobra application.`,

	Run: func(cmd *cobra.Command, args []string) {

		apiUrl:= "/api/v2/users/myself"
		endpoint:= Endpoint(apiUrl)

		response, err := http.Get(endpoint);
		if err != nil {
			fmt.Print(err.Error())
			os.Exit(1)
		}

		responseData, err := ioutil.ReadAll(response.Body)

		if err != nil {
			log.Fatal(err)
		}
		fmt.Println(string(responseData))

	},
}

func init() {
	RootCmd.AddCommand(meCmd)
}
<|endoftext|>


<commit_before>package lzma2

import (
	"errors"
	"fmt"

	"github.com/uli-go/xz/lzma"
)

type chunkType byte

const (
	cEOS chunkType = iota
	cU
	cUD
	cL
	cLR
	cLRN
	cLRND
)

var chunkTypeStrings = [...]string{
	cEOS:  "EOS",
	cU:    "U",
	cUD:   "UD",
	cL:    "L",
	cLR:   "LR",
	cLRN:  "LRN",
	cLRND: "LRND",
}

func (c chunkType) String() string {
	if !(cEOS <= c && c <= cLRND) {
		return "unknown"
	}
	return chunkTypeStrings[c]
}

type chunkHeader struct {
	ctype    chunkType
	unpacked uint32
	packed   uint16
	props    lzma.Properties
}

const (
	hEOS  = 0
	hUD   = 1
	hU    = 2
	hL    = 1 << 7
	hLR   = 1<<7 | 1<<5
	hLRN  = 1<<7 | 1<<6
	hLRND = 1<<7 | 1<<6 | 1<<5
)

var errHeaderByte = errors.New("unsupported chunk header byte")

func headerChunkType(h byte) (c chunkType, err error) {
	if h&hL == 0 {
		// no compression
		switch h {
		case hEOS:
			c = cEOS
		case hUD:
			c = cUD
		case hU:
			c = cU
		default:
			return 0, errHeaderByte
		}
		return
	}
	switch h & hLRND {
	case hL:
		c = cL
	case hLR:
		c = cLR
	case hLRN:
		c = cLRN
	case hLRND:
		c = cLRND
	default:
		return 0, errHeaderByte
	}
	return
}

func headerLen(c chunkType) int {
	switch c {
	case cEOS:
		return 1
	case cU, cUD:
		return 3
	case cL, cLR:
		return 5
	case cLRN, cLRND:
		return 6
	}
	panic(fmt.Sprintf("unsupported chunk type %d", c))
}
<commit_msg>lzma2: more work on read header function<commit_after>package lzma2

import (
	"errors"
	"fmt"
	"io"

	"github.com/uli-go/xz/lzma"
)

type chunkType byte

const (
	cEOS chunkType = iota
	cU
	cUD
	cL
	cLR
	cLRN
	cLRND
)

var chunkTypeStrings = [...]string{
	cEOS:  "EOS",
	cU:    "U",
	cUD:   "UD",
	cL:    "L",
	cLR:   "LR",
	cLRN:  "LRN",
	cLRND: "LRND",
}

func (c chunkType) String() string {
	if !(cEOS <= c && c <= cLRND) {
		return "unknown"
	}
	return chunkTypeStrings[c]
}

type chunkHeader struct {
	ctype    chunkType
	unpacked uint32
	packed   uint16
	props    lzma.Properties
}

const (
	hEOS  = 0
	hUD   = 1
	hU    = 2
	hL    = 1 << 7
	hLR   = 1<<7 | 1<<5
	hLRN  = 1<<7 | 1<<6
	hLRND = 1<<7 | 1<<6 | 1<<5
)

var errHeaderByte = errors.New("unsupported chunk header byte")

func headerChunkType(h byte) (c chunkType, err error) {
	if h&hL == 0 {
		// no compression
		switch h {
		case hEOS:
			c = cEOS
		case hUD:
			c = cUD
		case hU:
			c = cU
		default:
			return 0, errHeaderByte
		}
		return
	}
	switch h & hLRND {
	case hL:
		c = cL
	case hLR:
		c = cLR
	case hLRN:
		c = cLRN
	case hLRND:
		c = cLRND
	default:
		return 0, errHeaderByte
	}
	return
}

func headerLen(c chunkType) int {
	switch c {
	case cEOS:
		return 1
	case cU, cUD:
		return 3
	case cL, cLR:
		return 5
	case cLRN, cLRND:
		return 6
	}
	panic(fmt.Sprintf("unsupported chunk type %d", c))
}

func readChunkHeader(r io.Reader) (h *chunkHeader, err error) {
	p := make([]byte, 1, 6)
	if _, err = io.ReadFull(r, p); err != nil {
		return
	}
	if c, err = headerChunkType(p[0]); err != nil {
		return
	}
	p = p[:headerLen(c)]
	if _, err = ReadFull(r, p[1:]); err != nil {
		return
	}
	panic("TODO")
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
package mysql

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/admpub/errors"
	"github.com/admpub/nging/application/library/dbmanager/driver"
	"github.com/webx-top/db"
)

/*
mysqldump 参数说明：
-d 			结构(--no-data:不导出任何数据，只导出数据库表结构)
-t 			数据(--no-create-info:只导出数据，而不添加CREATE TABLE 语句)
-n 			(--no-create-db:只导出数据，而不添加CREATE DATABASE 语句）
-R 			(--routines:导出存储过程以及自定义函数)
-E 			(--events:导出事件)
--triggers 	(默认导出触发器，使用--skip-triggers屏蔽导出)
-B 			(--databases:导出数据库列表，单个库时可省略）
--tables 	表列表（单个表时可省略）
*/

var cleanRegExp = regexp.MustCompile(` AUTO_INCREMENT=[0-9]*\s*`)

// Export 导出SQL文件
func Export(cfg *driver.DbAuth, tables []string, structWriter, dataWriter interface{}, resetAutoIncrements ...bool) error {
	if len(tables) == 0 {
		return nil
	}
	log.Println(`Starting backup:`, tables)
	var (
		port, host         string
		resetAutoIncrement bool
	)
	if len(resetAutoIncrements) > 0 {
		resetAutoIncrement = resetAutoIncrements[0]
	}
	if p := strings.LastIndex(cfg.Host, `:`); p > 0 {
		host = cfg.Host[0:p]
		port = cfg.Host[p+1:]
	} else {
		host = cfg.Host
	}
	if len(port) == 0 {
		port = `3306`
	}
	args := []string{
		"--default-character-set=" + cfg.Charset,
		"--single-transaction=TRUE",
		"--column-statistics=0",
		"--set-gtid-purged=OFF",
		"--no-autocommit",
		//"--ignore-table="+cfg.Db+".details",
		//"--ignore-table="+cfg.Db+".large_table2"
		"--opt",
		"-d", //加上此参数代表只导出表结构，不导出数据
		"-h" + host,
		"-P" + port,
		"-u" + cfg.Username,
		"-p" + cfg.Password,
		cfg.Db,
	}
	clean := func(w io.Writer) {
		if c, y := w.(io.Closer); y {
			c.Close()
		}
	}
	var typeOptIndex int
	for index, value := range args {
		if value == `-d` {
			typeOptIndex = index
			break
		}
	}
	args = append(args, tables...)
	for index, writer := range []interface{}{structWriter, dataWriter} {
		if writer == nil {
			continue
		}
		if index > 0 {
			args[typeOptIndex] = `-t` //导出数据
		}
		cmd := exec.Command("mysqldump", args...)
		var (
			w        io.Writer
			err      error
			onFinish func() error
		)
		switch v := writer.(type) {
		case io.Writer:
			w = v
		case string:
			dir := filepath.Dir(v)
			if _, err := os.Stat(dir); os.IsNotExist(err) {
				err = os.MkdirAll(dir, os.ModePerm)
				if err != nil {
					return fmt.Errorf(`Failed to backup: %v`, err)
				}
			}
			w, err = os.Create(v)
			if err != nil {
				return fmt.Errorf(`Failed to backup: %v`, err)
			}
			onFinish = func() error {
				if index > 0 {
					return nil
				}
				if resetAutoIncrement {
					return ResetAutoIncrement(v)
				}
				return nil
			}
		default:
			return errors.Wrapf(db.ErrUnsupported, `SQL Writer Error: %T`, v)
		}
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if err := cmd.Start(); err != nil {
			stdout.Close()
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if _, err := io.Copy(w, stdout); err != nil {
			stdout.Close()
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		clean(w)
		stdout.Close()
		if onFinish != nil {
			if err = onFinish(); err != nil {
				return err
			}
		}
	}
	return nil
}

// ResetAutoIncrement 重置AUTO_INCREMENT值为0
func ResetAutoIncrement(sqlStructFile string) error {
	b, err := ioutil.ReadFile(sqlStructFile)
	if err != nil {
		return err
	}
	b = cleanRegExp.ReplaceAll(b, []byte(` `))
	return ioutil.WriteFile(sqlStructFile, b, 0666)
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
package mysql

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/admpub/errors"
	"github.com/admpub/nging/application/library/dbmanager/driver"
	"github.com/webx-top/db"
)

/*
mysqldump 参数说明：
-d 			结构(--no-data:不导出任何数据，只导出数据库表结构)
-t 			数据(--no-create-info:只导出数据，而不添加CREATE TABLE 语句)
-n 			(--no-create-db:只导出数据，而不添加CREATE DATABASE 语句）
-R 			(--routines:导出存储过程以及自定义函数)
-E 			(--events:导出事件)
--triggers 	(默认导出触发器，使用--skip-triggers屏蔽导出)
-B 			(--databases:导出数据库列表，单个库时可省略）
--tables 	表列表（单个表时可省略）
*/

var cleanRegExp = regexp.MustCompile(` AUTO_INCREMENT=[0-9]*\s*`)

// Export 导出SQL文件
func Export(cfg *driver.DbAuth, tables []string, structWriter, dataWriter interface{}, resetAutoIncrements ...bool) error {
	if len(tables) == 0 {
		return nil
	}
	log.Println(`Starting backup:`, tables)
	var (
		port, host         string
		resetAutoIncrement bool
	)
	if len(resetAutoIncrements) > 0 {
		resetAutoIncrement = resetAutoIncrements[0]
	}
	if p := strings.LastIndex(cfg.Host, `:`); p > 0 {
		host = cfg.Host[0:p]
		port = cfg.Host[p+1:]
	} else {
		host = cfg.Host
	}
	if len(port) == 0 {
		port = `3306`
	}
	args := []string{
		"--default-character-set=" + cfg.Charset,
		"--single-transaction=TRUE",
		"--column-statistics=0",
		"--set-gtid-purged=OFF",
		"--no-autocommit",
		//"--ignore-table="+cfg.Db+".details",
		//"--ignore-table="+cfg.Db+".large_table2"
		"--opt",
		"-d", //加上此参数代表只导出表结构，不导出数据
		"-h" + host,
		"-P" + port,
		"-u" + cfg.Username,
		"-p" + cfg.Password,
		cfg.Db,
		//"--result-file=/root/backup.sql",
	}
	clean := func(w io.Writer) {
		if c, y := w.(io.Closer); y {
			c.Close()
		}
	}
	var typeOptIndex int
	for index, value := range args {
		if value == `-d` {
			typeOptIndex = index
			break
		}
	}
	args = append(args, tables...)
	for index, writer := range []interface{}{structWriter, dataWriter} {
		if writer == nil {
			continue
		}
		if index > 0 {
			args[typeOptIndex] = `-t` //导出数据
		}
		cmd := exec.Command("mysqldump", args...)
		var (
			w        io.Writer
			err      error
			onFinish func() error
		)
		switch v := writer.(type) {
		case io.Writer:
			w = v
		case string:
			dir := filepath.Dir(v)
			if _, err := os.Stat(dir); os.IsNotExist(err) {
				err = os.MkdirAll(dir, os.ModePerm)
				if err != nil {
					return fmt.Errorf(`Failed to backup: %v`, err)
				}
			}
			w, err = os.Create(v)
			if err != nil {
				return fmt.Errorf(`Failed to backup: %v`, err)
			}
			onFinish = func() error {
				if index > 0 {
					return nil
				}
				if resetAutoIncrement {
					return ResetAutoIncrement(v)
				}
				return nil
			}
		default:
			return errors.Wrapf(db.ErrUnsupported, `SQL Writer Error: %T`, v)
		}
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if err := cmd.Start(); err != nil {
			stdout.Close()
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if _, err := io.Copy(w, stdout); err != nil {
			stdout.Close()
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		clean(w)
		stdout.Close()
		if onFinish != nil {
			if err = onFinish(); err != nil {
				return err
			}
		}
	}
	return nil
}

// ResetAutoIncrement 重置AUTO_INCREMENT值为0
func ResetAutoIncrement(sqlStructFile string) error {
	b, err := ioutil.ReadFile(sqlStructFile)
	if err != nil {
		return err
	}
	b = cleanRegExp.ReplaceAll(b, []byte(` `))
	return ioutil.WriteFile(sqlStructFile, b, 0666)
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package mysql

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/admpub/archiver"
	"github.com/admpub/errors"
	loga "github.com/admpub/log"
	"github.com/admpub/nging/application/handler"
	"github.com/admpub/nging/application/library/cron"
	"github.com/admpub/nging/application/library/dbmanager/driver"
	"github.com/webx-top/com"
	"github.com/webx-top/db"
	"github.com/webx-top/echo"
)

/*
mysqldump 参数说明：
-d 			结构(--no-data:不导出任何数据，只导出数据库表结构)
-t 			数据(--no-create-info:只导出数据，而不添加CREATE TABLE 语句)
-n 			(--no-create-db:只导出数据，而不添加CREATE DATABASE 语句）
-R 			(--routines:导出存储过程以及自定义函数)
-E 			(--events:导出事件)
--triggers 	(默认导出触发器，使用--skip-triggers屏蔽导出)
-B 			(--databases:导出数据库列表，单个库时可省略）
--tables 	表列表（单个表时可省略）
*/

var (
	cleanRegExp = regexp.MustCompile(` AUTO_INCREMENT=[0-9]*\s*`)
	backgrounds = sync.Map{} //后台导入导出任务
	// SQLTempDir sql文件缓存目录获取函数(用于导入导出SQL)
	SQLTempDir = os.TempDir
)

const (
	OpExport = `export`
	OpImport = `export`
)

type FileInfos []*FileInfo

type FileInfo struct {
	Start      time.Time
	End        time.Time
	Elapsed    time.Duration
	Path       string
	Size       int64
	Compressed bool
	Error      string
}

func TempDir(op string) string {
	dir := filepath.Join(SQLTempDir(), `dbmanager/cache`, op)
	if !com.IsDir(dir) {
		err := os.MkdirAll(dir, os.ModePerm)
		if err != nil {
			loga.Error(err)
		}
	}
	return dir
}

// Export 导出SQL文件
func Export(ctx context.Context, cfg *driver.DbAuth, tables []string, structWriter, dataWriter interface{}, resetAutoIncrements ...bool) error {
	if len(tables) == 0 {
		return errors.New(`No table selected for export`)
	}
	log.Println(`Starting backup:`, tables)
	var (
		port, host         string
		resetAutoIncrement bool
	)
	if len(resetAutoIncrements) > 0 {
		resetAutoIncrement = resetAutoIncrements[0]
	}
	if p := strings.LastIndex(cfg.Host, `:`); p > 0 {
		host = cfg.Host[0:p]
		port = cfg.Host[p+1:]
	} else {
		host = cfg.Host
	}
	if len(port) == 0 {
		port = `3306`
	}
	if !com.InSlice(cfg.Charset, Charsets) {
		return errors.New(`字符集charset值无效`)
	}
	args := []string{
		"--default-character-set=" + cfg.Charset,
		"--single-transaction",
		//"--column-statistics=0",//低版本不支持
		"--set-gtid-purged=OFF",
		"--no-autocommit",
		//"--ignore-table="+cfg.Db+".details",
		//"--ignore-table="+cfg.Db+".large_table2",
		"--opt",
		"-d", //加上此参数代表只导出表结构，不导出数据
		"-h" + host,
		"-P" + port,
		"-u" + cfg.Username,
		"-p" + cfg.Password,
		cfg.Db,
		//"--result-file=/root/backup.sql",
	}
	clean := func(w io.Writer) {
		if c, y := w.(io.Closer); y {
			c.Close()
		}
	}
	var typeOptIndex int
	for index, value := range args {
		if value == `-d` {
			typeOptIndex = index
			break
		}
	}
	//args = append(args, `--tables`)
	args = append(args, tables...)
	rec := cron.NewCmdRec(1000)
	for index, writer := range []interface{}{structWriter, dataWriter} {
		if writer == nil {
			continue
		}
		if index > 0 {
			args[typeOptIndex] = `-t` //导出数据
		}
		//log.Println(`mysqldump`, strings.Join(args, ` `))
		cmd := exec.CommandContext(ctx, "mysqldump", args...)
		cmd.Stderr = rec
		var (
			w        io.Writer
			err      error
			onFinish func() error
		)
		switch v := writer.(type) {
		case io.Writer:
			w = v
		case string:
			dir := filepath.Dir(v)
			if _, err := os.Stat(dir); os.IsNotExist(err) {
				err = os.MkdirAll(dir, os.ModePerm)
				if err != nil {
					return fmt.Errorf(`Failed to backup: %v`, err)
				}
			}
			w, err = os.Create(v)
			if err != nil {
				return fmt.Errorf(`Failed to backup: %v`, err)
			}
			onFinish = func() error {
				if index > 0 {
					return nil
				}
				if resetAutoIncrement {
					return ResetAutoIncrement(v)
				}
				return nil
			}
		default:
			return errors.Wrapf(db.ErrUnsupported, `SQL Writer Error: %T`, v)
		}
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if err := cmd.Start(); err != nil {
			stdout.Close()
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if _, err := io.Copy(w, stdout); err != nil {
			stdout.Close()
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if err := cmd.Wait(); err != nil {
			stdout.Close()
			clean(w)
			return errors.New(err.Error() + `: ` + rec.String())
		}
		clean(w)
		stdout.Close()
		if onFinish != nil {
			if err = onFinish(); err != nil {
				return err
			}
		}
	}
	return nil
}

// ResetAutoIncrement 重置AUTO_INCREMENT值为0
func ResetAutoIncrement(sqlStructFile string) error {
	b, err := ioutil.ReadFile(sqlStructFile)
	if err != nil {
		return err
	}
	b = cleanRegExp.ReplaceAll(b, []byte(` `))
	return ioutil.WriteFile(sqlStructFile, b, 0666)
}

func (m *mySQL) getCharsetList() []string {
	e, _ := m.getCharsets()
	cs := make([]string, 0)
	for k := range e {
		cs = append(cs, k)
	}
	sort.Strings(cs)
	return cs
}

func (m *mySQL) Export() error {
	//fmt.Printf("%#v\n", m.getCharsetList())
	var err error
	if m.IsPost() {
		var tables []string
		if len(m.dbName) == 0 {
			m.fail(m.T(`请选择数据库`))
			return m.returnTo(m.GenURL(`listDb`))
		}
		if m.Formx(`all`).Bool() {
			var ok bool
			tables, ok = m.Get(`tableList`).([]string)
			if !ok {
				tables, err = m.getTables()
				if err != nil {
					m.fail(err.Error())
					return m.returnTo(m.GenURL(`export`))
				}
			}
		} else {
			tables = m.FormValues(`table`)
			if len(tables) == 1 && len(tables[0]) > 0 {
				tables = strings.Split(tables[0], `,`)
			}
			views := m.FormValues(`view`)
			if len(views) == 1 && len(views[0]) > 0 {
				views = append(tables, strings.Split(views[0], `,`)...)
			}
			if len(views) > 0 {
				tables = append(tables, views...)
			}
		}
		output := m.Form(`output`)
		types := m.FormValues(`type`)
		cacheKey := com.Md5(com.Dump([]interface{}{tables, output, types}, false))

		var exports map[string]*FileInfos
		if old, exists := backgrounds.Load(Export); exists {
			exports = old.(map[string]*FileInfos)
		} else {
			exports = map[string]*FileInfos{}
		}
		if _, ok := exports[cacheKey]; !ok {
			return errors.New(m.T(`任务正在后台处理中，请稍候...`))
		}

		var (
			structWriter, dataWriter interface{}
			sqlFiles                 []string
			async                    bool
			fileInfos                = &FileInfos{}
		)
		exports[cacheKey] = fileInfos
		nowTime := com.String(time.Now().Unix())
		saveDir := TempDir(`export`)
		switch output {
		case `down`:
			m.Response().Header().Set(echo.HeaderContentDisposition, fmt.Sprintf("attachment; filename=%q", m.dbName+"-sql-"+nowTime+".sql"))
			fallthrough
		case `open`:
			if com.InSlice(`struct`, types) {
				structWriter = m.Response()
			}
			if com.InSlice(`data`, types) {
				dataWriter = m.Response()
			}
		default:
			async = true
			if com.InSlice(`struct`, types) {
				structFile := filepath.Join(saveDir, m.dbName+`-struct-`+nowTime+`.sql`)
				sqlFiles = append(sqlFiles, structFile)
				structWriter = structFile
				fi := &FileInfo{
					Start: time.Now(),
					Path:  structFile,
				}
				*fileInfos = append(*fileInfos, fi)
			}
			if com.InSlice(`data`, types) {
				dataFile := filepath.Join(saveDir, m.dbName+`-data-`+nowTime+`.sql`)
				sqlFiles = append(sqlFiles, dataFile)
				dataWriter = dataFile
				fi := &FileInfo{
					Start: time.Now(),
					Path:  dataFile,
				}
				*fileInfos = append(*fileInfos, fi)
			}
		}
		cfg := *m.DbAuth
		cfg.Db = m.dbName
		worker := func(ctx context.Context) error {
			defer func() {
				if r := recover(); r != nil {
					err = fmt.Errorf(`%v`, r)
				}
			}()
			err = Export(ctx, &cfg, tables, structWriter, dataWriter, true)
			if err != nil {
				loga.Error(err)
				return err
			}
			if len(sqlFiles) > 0 {
				now := time.Now()
				for _, fi := range *fileInfos {
					fi.End = now
					fi.Size, err = com.FileSize(fi.Path)
					if err != nil {
						fi.Error = err.Error()
					}
					fi.Elapsed = fi.End.Sub(fi.Start)
				}
				zipFile := filepath.Join(saveDir, m.dbName+"-sql-"+nowTime+".zip")
				fi := &FileInfo{
					Start:      now,
					Path:       zipFile,
					Compressed: true,
				}
				err = archiver.Zip.Make(zipFile, sqlFiles)
				if err != nil {
					loga.Error(err)
					return err
				}
				for _, sqlFile := range sqlFiles {
					os.Remove(sqlFile)
				}
				fi.Size, err = com.FileSize(zipFile)
				if err != nil {
					fi.Error = err.Error()
				}
				fi.End = time.Now()
				fi.Elapsed = fi.End.Sub(fi.Start)
				*fileInfos = append(*fileInfos, fi)
				ioutil.WriteFile(zipFile+`.txt`, com.Str2bytes(com.Dump(fileInfos, false)), os.ModePerm)
				if _, ok := exports[cacheKey]; ok {
					delete(exports, cacheKey)
				}
				/*
					fp, err := os.Open(zipFile)
					if err != nil {
						loga.Error(err)
						return err
					}
					defer func() {
						fp.Close()
						os.Remove(zipFile)
					}()
					return m.Attachment(fp, filepath.Base(zipFile))
				*/
			}
			return nil
		}
		if !async {
			ctx, cancel := context.WithCancel(m.StdContext())
			done := make(chan error)
			clientGone := m.Response().(http.CloseNotifier).CloseNotify()
			go func() {
				for {
					select {
					case <-clientGone:
						cancel()
						return
					case <-done:
						return
					}
				}
			}()
			err = worker(ctx)
			done <- err
			return err
		}
		data := m.Data()
		data.SetInfo(m.T(`任务已经在后台成功启动`))
		data.SetURL(handler.URLFor(`/download/file?path=dbmanager/cache/export`))
		go worker(m)
		exports[cacheKey] = fileInfos
		backgrounds.Store(OpExport, exports)
		return m.JSON(data)
	}
	return m.Redirect(m.GenURL(`listTable`, m.dbName))
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package mysql

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/admpub/archiver"
	"github.com/admpub/errors"
	loga "github.com/admpub/log"
	"github.com/admpub/nging/application/handler"
	"github.com/admpub/nging/application/library/cron"
	"github.com/admpub/nging/application/library/dbmanager/driver"
	"github.com/webx-top/com"
	"github.com/webx-top/db"
	"github.com/webx-top/echo"
)

/*
mysqldump 参数说明：
-d 			结构(--no-data:不导出任何数据，只导出数据库表结构)
-t 			数据(--no-create-info:只导出数据，而不添加CREATE TABLE 语句)
-n 			(--no-create-db:只导出数据，而不添加CREATE DATABASE 语句）
-R 			(--routines:导出存储过程以及自定义函数)
-E 			(--events:导出事件)
--triggers 	(默认导出触发器，使用--skip-triggers屏蔽导出)
-B 			(--databases:导出数据库列表，单个库时可省略）
--tables 	表列表（单个表时可省略）
*/

var (
	cleanRegExp = regexp.MustCompile(` AUTO_INCREMENT=[0-9]*\s*`)
	backgrounds = sync.Map{} //后台导入导出任务
	// SQLTempDir sql文件缓存目录获取函数(用于导入导出SQL)
	SQLTempDir = os.TempDir
)

const (
	OpExport = `export`
	OpImport = `export`
)

type FileInfos []*FileInfo

type FileInfo struct {
	Start      time.Time
	End        time.Time
	Elapsed    time.Duration
	Path       string
	Size       int64
	Compressed bool
	Error      string
}

func TempDir(op string) string {
	dir := filepath.Join(SQLTempDir(), `dbmanager/cache`, op)
	if !com.IsDir(dir) {
		err := os.MkdirAll(dir, os.ModePerm)
		if err != nil {
			loga.Error(err)
		}
	}
	return dir
}

// Export 导出SQL文件
func Export(ctx context.Context, cfg *driver.DbAuth, tables []string, structWriter, dataWriter interface{}, resetAutoIncrements ...bool) error {
	if len(tables) == 0 {
		return errors.New(`No table selected for export`)
	}
	log.Println(`Starting backup:`, tables)
	var (
		port, host         string
		resetAutoIncrement bool
	)
	if len(resetAutoIncrements) > 0 {
		resetAutoIncrement = resetAutoIncrements[0]
	}
	if p := strings.LastIndex(cfg.Host, `:`); p > 0 {
		host = cfg.Host[0:p]
		port = cfg.Host[p+1:]
	} else {
		host = cfg.Host
	}
	if len(port) == 0 {
		port = `3306`
	}
	if !com.InSlice(cfg.Charset, Charsets) {
		return errors.New(`字符集charset值无效`)
	}
	args := []string{
		"--default-character-set=" + cfg.Charset,
		"--single-transaction",
		//"--column-statistics=0",//低版本不支持
		"--set-gtid-purged=OFF",
		"--no-autocommit",
		//"--ignore-table="+cfg.Db+".details",
		//"--ignore-table="+cfg.Db+".large_table2",
		"--opt",
		"-d", //加上此参数代表只导出表结构，不导出数据
		"-h" + host,
		"-P" + port,
		"-u" + cfg.Username,
		"-p" + cfg.Password,
		cfg.Db,
		//"--result-file=/root/backup.sql",
	}
	clean := func(w io.Writer) {
		if c, y := w.(io.Closer); y {
			c.Close()
		}
	}
	var typeOptIndex int
	for index, value := range args {
		if value == `-d` {
			typeOptIndex = index
			break
		}
	}
	//args = append(args, `--tables`)
	args = append(args, tables...)
	rec := cron.NewCmdRec(1000)
	for index, writer := range []interface{}{structWriter, dataWriter} {
		if writer == nil {
			continue
		}
		if index > 0 {
			args[typeOptIndex] = `-t` //导出数据
		}
		//log.Println(`mysqldump`, strings.Join(args, ` `))
		cmd := exec.CommandContext(ctx, "mysqldump", args...)
		cmd.Stderr = rec
		var (
			w        io.Writer
			err      error
			onFinish func() error
		)
		switch v := writer.(type) {
		case io.Writer:
			w = v
		case string:
			dir := filepath.Dir(v)
			if _, err := os.Stat(dir); os.IsNotExist(err) {
				err = os.MkdirAll(dir, os.ModePerm)
				if err != nil {
					return fmt.Errorf(`Failed to backup: %v`, err)
				}
			}
			w, err = os.Create(v)
			if err != nil {
				return fmt.Errorf(`Failed to backup: %v`, err)
			}
			onFinish = func() error {
				if index > 0 {
					return nil
				}
				if resetAutoIncrement {
					return ResetAutoIncrement(v)
				}
				return nil
			}
		default:
			return errors.Wrapf(db.ErrUnsupported, `SQL Writer Error: %T`, v)
		}
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if err := cmd.Start(); err != nil {
			stdout.Close()
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if _, err := io.Copy(w, stdout); err != nil {
			stdout.Close()
			clean(w)
			return fmt.Errorf(`Failed to backup: %v`, err)
		}
		if err := cmd.Wait(); err != nil {
			stdout.Close()
			clean(w)
			return errors.New(err.Error() + `: ` + rec.String())
		}
		clean(w)
		stdout.Close()
		if onFinish != nil {
			if err = onFinish(); err != nil {
				return err
			}
		}
	}
	return nil
}

// ResetAutoIncrement 重置AUTO_INCREMENT值为0
func ResetAutoIncrement(sqlStructFile string) error {
	b, err := ioutil.ReadFile(sqlStructFile)
	if err != nil {
		return err
	}
	b = cleanRegExp.ReplaceAll(b, []byte(` `))
	return ioutil.WriteFile(sqlStructFile, b, 0666)
}

func (m *mySQL) getCharsetList() []string {
	e, _ := m.getCharsets()
	cs := make([]string, 0)
	for k := range e {
		cs = append(cs, k)
	}
	sort.Strings(cs)
	return cs
}

func (m *mySQL) Export() error {
	//fmt.Printf("%#v\n", m.getCharsetList())
	var err error
	if m.IsPost() {
		var tables []string
		if len(m.dbName) == 0 {
			m.fail(m.T(`请选择数据库`))
			return m.returnTo(m.GenURL(`listDb`))
		}
		if m.Formx(`all`).Bool() {
			var ok bool
			tables, ok = m.Get(`tableList`).([]string)
			if !ok {
				tables, err = m.getTables()
				if err != nil {
					m.fail(err.Error())
					return m.returnTo(m.GenURL(`export`))
				}
			}
		} else {
			tables = m.FormValues(`table`)
			if len(tables) == 1 && len(tables[0]) > 0 {
				tables = strings.Split(tables[0], `,`)
			}
			views := m.FormValues(`view`)
			if len(views) == 1 && len(views[0]) > 0 {
				views = append(tables, strings.Split(views[0], `,`)...)
			}
			if len(views) > 0 {
				tables = append(tables, views...)
			}
		}
		output := m.Form(`output`)
		types := m.FormValues(`type`)
		cacheKey := com.Md5(com.Dump([]interface{}{tables, output, types}, false))

		var exports map[string]*FileInfos
		if old, exists := backgrounds.Load(OpExport); exists {
			exports = old.(map[string]*FileInfos)
		} else {
			exports = map[string]*FileInfos{}
		}
		if _, ok := exports[cacheKey]; !ok {
			return errors.New(m.T(`任务正在后台处理中，请稍候...`))
		}

		var (
			structWriter, dataWriter interface{}
			sqlFiles                 []string
			async                    bool
			fileInfos                = &FileInfos{}
		)
		exports[cacheKey] = fileInfos
		nowTime := com.String(time.Now().Unix())
		saveDir := TempDir(`export`)
		switch output {
		case `down`:
			m.Response().Header().Set(echo.HeaderContentDisposition, fmt.Sprintf("attachment; filename=%q", m.dbName+"-sql-"+nowTime+".sql"))
			fallthrough
		case `open`:
			if com.InSlice(`struct`, types) {
				structWriter = m.Response()
			}
			if com.InSlice(`data`, types) {
				dataWriter = m.Response()
			}
		default:
			async = true
			if com.InSlice(`struct`, types) {
				structFile := filepath.Join(saveDir, m.dbName+`-struct-`+nowTime+`.sql`)
				sqlFiles = append(sqlFiles, structFile)
				structWriter = structFile
				fi := &FileInfo{
					Start: time.Now(),
					Path:  structFile,
				}
				*fileInfos = append(*fileInfos, fi)
			}
			if com.InSlice(`data`, types) {
				dataFile := filepath.Join(saveDir, m.dbName+`-data-`+nowTime+`.sql`)
				sqlFiles = append(sqlFiles, dataFile)
				dataWriter = dataFile
				fi := &FileInfo{
					Start: time.Now(),
					Path:  dataFile,
				}
				*fileInfos = append(*fileInfos, fi)
			}
		}
		cfg := *m.DbAuth
		cfg.Db = m.dbName
		worker := func(ctx context.Context) error {
			defer func() {
				if r := recover(); r != nil {
					err = fmt.Errorf(`%v`, r)
				}
			}()
			err = Export(ctx, &cfg, tables, structWriter, dataWriter, true)
			if err != nil {
				loga.Error(err)
				return err
			}
			if len(sqlFiles) > 0 {
				now := time.Now()
				for _, fi := range *fileInfos {
					fi.End = now
					fi.Size, err = com.FileSize(fi.Path)
					if err != nil {
						fi.Error = err.Error()
					}
					fi.Elapsed = fi.End.Sub(fi.Start)
				}
				zipFile := filepath.Join(saveDir, m.dbName+"-sql-"+nowTime+".zip")
				fi := &FileInfo{
					Start:      now,
					Path:       zipFile,
					Compressed: true,
				}
				err = archiver.Zip.Make(zipFile, sqlFiles)
				if err != nil {
					loga.Error(err)
					return err
				}
				for _, sqlFile := range sqlFiles {
					os.Remove(sqlFile)
				}
				fi.Size, err = com.FileSize(zipFile)
				if err != nil {
					fi.Error = err.Error()
				}
				fi.End = time.Now()
				fi.Elapsed = fi.End.Sub(fi.Start)
				*fileInfos = append(*fileInfos, fi)
				ioutil.WriteFile(zipFile+`.txt`, com.Str2bytes(com.Dump(fileInfos, false)), os.ModePerm)
				if _, ok := exports[cacheKey]; ok {
					delete(exports, cacheKey)
				}
				/*
					fp, err := os.Open(zipFile)
					if err != nil {
						loga.Error(err)
						return err
					}
					defer func() {
						fp.Close()
						os.Remove(zipFile)
					}()
					return m.Attachment(fp, filepath.Base(zipFile))
				*/
			}
			return nil
		}
		if !async {
			ctx, cancel := context.WithCancel(m.StdContext())
			done := make(chan error)
			clientGone := m.Response().(http.CloseNotifier).CloseNotify()
			go func() {
				for {
					select {
					case <-clientGone:
						cancel()
						return
					case <-done:
						return
					}
				}
			}()
			err = worker(ctx)
			done <- err
			return err
		}
		data := m.Data()
		data.SetInfo(m.T(`任务已经在后台成功启动`))
		data.SetURL(handler.URLFor(`/download/file?path=dbmanager/cache/export`))
		go worker(m)
		exports[cacheKey] = fileInfos
		backgrounds.Store(OpExport, exports)
		return m.JSON(data)
	}
	return m.Redirect(m.GenURL(`listTable`, m.dbName))
}
<|endoftext|>