<commit_before>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

// Command tipgodoc is the beginning of the new tip.golang.org server,
// serving the latest HEAD straight from the Git oven.
package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"
)

const (
	repoURL      = "https://go.googlesource.com/"
	metaURL      = "https://go.googlesource.com/?b=master&format=JSON"
	startTimeout = 10 * time.Minute
)

func main() {
	const k = "TIP_BUILDER"
	var b Builder
	switch os.Getenv(k) {
	case "godoc":
		b = godocBuilder{}
	case "talks":
		b = talksBuilder{}
	default:
		log.Fatalf("Unknown %v value: %q", k, os.Getenv(k))
	}

	p := &Proxy{builder: b}
	go p.run()
	http.Handle("/", httpsOnlyHandler{p})
	http.HandleFunc("/_ah/health", p.serveHealthCheck)

	log.Print("Starting up")

	if err := http.ListenAndServe(":8080", nil); err != nil {
		p.stop()
		log.Fatal(err)
	}
}

// Proxy implements the tip.golang.org server: a reverse-proxy
// that builds and runs godoc instances showing the latest docs.
type Proxy struct {
	builder Builder

	mu       sync.Mutex // protects the followin'
	proxy    http.Handler
	cur      string    // signature of gorepo+toolsrepo
	cmd      *exec.Cmd // live godoc instance, or nil for none
	side     string
	hostport string // host and port of the live instance
	err      error
}

type Builder interface {
	Signature(heads map[string]string) string
	Init(dir, hostport string, heads map[string]string) (*exec.Cmd, error)
	HealthCheck(hostport string) error
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/_tipstatus" {
		p.serveStatus(w, r)
		return
	}
	p.mu.Lock()
	proxy := p.proxy
	err := p.err
	p.mu.Unlock()
	if proxy == nil {
		s := "starting up"
		if err != nil {
			s = err.Error()
		}
		http.Error(w, s, http.StatusInternalServerError)
		return
	}
	proxy.ServeHTTP(w, r)
}

func (p *Proxy) serveStatus(w http.ResponseWriter, r *http.Request) {
	p.mu.Lock()
	defer p.mu.Unlock()
	fmt.Fprintf(w, "side=%v\ncurrent=%v\nerror=%v\n", p.side, p.cur, p.err)
}

func (p *Proxy) serveHealthCheck(w http.ResponseWriter, r *http.Request) {
	p.mu.Lock()
	defer p.mu.Unlock()
	// NOTE: Status 502, 503, 504 are the only status codes that signify an unhealthy app.
	// So long as this handler returns one of those codes, this instance will not be sent any requests.
	if p.proxy == nil {
		log.Printf("Health check: not ready")
		http.Error(w, "Not ready", http.StatusServiceUnavailable)
		return
	}

	if err := p.builder.HealthCheck(p.hostport); err != nil {
		log.Printf("Health check failed: %v", err)
		http.Error(w, "Health check failed", http.StatusServiceUnavailable)
		return
	}
	io.WriteString(w, "ok")
}

// run runs in its own goroutine.
func (p *Proxy) run() {
	p.side = "a"
	for {
		p.poll()
		time.Sleep(30 * time.Second)
	}
}

func (p *Proxy) stop() {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.cmd != nil {
		p.cmd.Process.Kill()
	}
}

// poll runs from the run loop goroutine.
func (p *Proxy) poll() {
	heads := gerritMetaMap()
	if heads == nil {
		return
	}

	sig := p.builder.Signature(heads)

	p.mu.Lock()
	changes := sig != p.cur
	curSide := p.side
	p.cur = sig
	p.mu.Unlock()

	if !changes {
		return
	}

	newSide := "b"
	if curSide == "b" {
		newSide = "a"
	}

	dir := filepath.Join(os.TempDir(), "tip", newSide)
	if err := os.MkdirAll(dir, 0755); err != nil {
		p.err = err
		return
	}
	hostport := "localhost:8081"
	if newSide == "b" {
		hostport = "localhost:8082"
	}
	cmd, err := p.builder.Init(dir, hostport, heads)
	if err != nil {
		err = fmt.Errorf("builder.Init: %v", err)
	} else {
		go func() {
			// TODO(adg,bradfitz): be smarter about dead processes
			if err := cmd.Wait(); err != nil {
				log.Printf("process in %v exited: %v", dir, err)
			}
		}()
		err = waitReady(p.builder, hostport)
		if err != nil {
			cmd.Process.Kill()
			err = fmt.Errorf("waitReady: %v", err)
		}
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if err != nil {
		log.Println(err)
		p.err = err
		return
	}

	u, err := url.Parse(fmt.Sprintf("http://%v/", hostport))
	if err != nil {
		err = fmt.Errorf("parsing hostport: %v", err)
		log.Println(err)
		p.err = err
		return
	}
	p.proxy = httputil.NewSingleHostReverseProxy(u)
	p.side = newSide
	p.hostport = hostport
	if p.cmd != nil {
		p.cmd.Process.Kill()
	}
	p.cmd = cmd
}

func waitReady(b Builder, hostport string) error {
	var err error
	deadline := time.Now().Add(startTimeout)
	for time.Now().Before(deadline) {
		if err = b.HealthCheck(hostport); err == nil {
			return nil
		}
		time.Sleep(time.Second)
	}
	return fmt.Errorf("timed out waiting for process at %v: %v", hostport, err)
}

func runErr(cmd *exec.Cmd) error {
	out, err := cmd.CombinedOutput()
	if err != nil {
		if len(out) == 0 {
			return err
		}
		return fmt.Errorf("%s\n%v", out, err)
	}
	return nil
}

func checkout(repo, hash, path string) error {
	// Clone git repo if it doesn't exist.
	if _, err := os.Stat(filepath.Join(path, ".git")); os.IsNotExist(err) {
		if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
			return fmt.Errorf("mkdir: %v", err)
		}
		if err := runErr(exec.Command("git", "clone", repo, path)); err != nil {
			return fmt.Errorf("clone: %v", err)
		}
	} else if err != nil {
		return fmt.Errorf("stat .git: %v", err)
	}

	// Pull down changes and update to hash.
	cmd := exec.Command("git", "fetch")
	cmd.Dir = path
	if err := runErr(cmd); err != nil {
		return fmt.Errorf("fetch: %v", err)
	}
	cmd = exec.Command("git", "reset", "--hard", hash)
	cmd.Dir = path
	if err := runErr(cmd); err != nil {
		return fmt.Errorf("reset: %v", err)
	}
	cmd = exec.Command("git", "clean", "-d", "-f", "-x")
	cmd.Dir = path
	if err := runErr(cmd); err != nil {
		return fmt.Errorf("clean: %v", err)
	}
	return nil
}

// gerritMetaMap returns the map from repo name (e.g. "go") to its
// latest master hash.
// The returned map is nil on any transient error.
func gerritMetaMap() map[string]string {
	res, err := http.Get(metaURL)
	if err != nil {
		return nil
	}
	defer res.Body.Close()
	defer io.Copy(ioutil.Discard, res.Body) // ensure EOF for keep-alive
	if res.StatusCode != 200 {
		return nil
	}
	var meta map[string]struct {
		Branches map[string]string
	}
	br := bufio.NewReader(res.Body)
	// For security reasons or something, this URL starts with ")]}'\n" before
	// the JSON object. So ignore that.
	// Shawn Pearce says it's guaranteed to always be just one line, ending in '\n'.
	for {
		b, err := br.ReadByte()
		if err != nil {
			return nil
		}
		if b == '\n' {
			break
		}
	}
	if err := json.NewDecoder(br).Decode(&meta); err != nil {
		log.Printf("JSON decoding error from %v: %s", metaURL, err)
		return nil
	}
	m := map[string]string{}
	for repo, v := range meta {
		if master, ok := v.Branches["master"]; ok {
			m[repo] = master
		}
	}
	return m
}

func getOK(url string) (body []byte, err error) {
	res, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	body, err = ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return nil, err
	}
	if res.StatusCode != http.StatusOK {
		return nil, errors.New(res.Status)
	}
	return body, nil
}

// httpsOnlyHandler redirects requests to "http://example.com/foo?bar"
// to "https://example.com/foo?bar"
type httpsOnlyHandler struct {
	h http.Handler
}

func (h httpsOnlyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Header.Get("X-Appengine-Https") == "off" {
		r.URL.Scheme = "https"
		r.URL.Host = r.Host
		http.Redirect(w, r, r.URL.String(), http.StatusFound)
		return
	}
	h.h.ServeHTTP(w, r)
}
<commit_msg>cmd/tip: set Strict-Transport-Security header in production<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

// Command tipgodoc is the beginning of the new tip.golang.org server,
// serving the latest HEAD straight from the Git oven.
package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"
)

const (
	repoURL      = "https://go.googlesource.com/"
	metaURL      = "https://go.googlesource.com/?b=master&format=JSON"
	startTimeout = 10 * time.Minute
)

func main() {
	const k = "TIP_BUILDER"
	var b Builder
	switch os.Getenv(k) {
	case "godoc":
		b = godocBuilder{}
	case "talks":
		b = talksBuilder{}
	default:
		log.Fatalf("Unknown %v value: %q", k, os.Getenv(k))
	}

	p := &Proxy{builder: b}
	go p.run()
	http.Handle("/", httpsOnlyHandler{p})
	http.HandleFunc("/_ah/health", p.serveHealthCheck)

	log.Print("Starting up")

	if err := http.ListenAndServe(":8080", nil); err != nil {
		p.stop()
		log.Fatal(err)
	}
}

// Proxy implements the tip.golang.org server: a reverse-proxy
// that builds and runs godoc instances showing the latest docs.
type Proxy struct {
	builder Builder

	mu       sync.Mutex // protects the followin'
	proxy    http.Handler
	cur      string    // signature of gorepo+toolsrepo
	cmd      *exec.Cmd // live godoc instance, or nil for none
	side     string
	hostport string // host and port of the live instance
	err      error
}

type Builder interface {
	Signature(heads map[string]string) string
	Init(dir, hostport string, heads map[string]string) (*exec.Cmd, error)
	HealthCheck(hostport string) error
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/_tipstatus" {
		p.serveStatus(w, r)
		return
	}
	p.mu.Lock()
	proxy := p.proxy
	err := p.err
	p.mu.Unlock()
	if proxy == nil {
		s := "starting up"
		if err != nil {
			s = err.Error()
		}
		http.Error(w, s, http.StatusInternalServerError)
		return
	}
	proxy.ServeHTTP(w, r)
}

func (p *Proxy) serveStatus(w http.ResponseWriter, r *http.Request) {
	p.mu.Lock()
	defer p.mu.Unlock()
	fmt.Fprintf(w, "side=%v\ncurrent=%v\nerror=%v\n", p.side, p.cur, p.err)
}

func (p *Proxy) serveHealthCheck(w http.ResponseWriter, r *http.Request) {
	p.mu.Lock()
	defer p.mu.Unlock()
	// NOTE: Status 502, 503, 504 are the only status codes that signify an unhealthy app.
	// So long as this handler returns one of those codes, this instance will not be sent any requests.
	if p.proxy == nil {
		log.Printf("Health check: not ready")
		http.Error(w, "Not ready", http.StatusServiceUnavailable)
		return
	}

	if err := p.builder.HealthCheck(p.hostport); err != nil {
		log.Printf("Health check failed: %v", err)
		http.Error(w, "Health check failed", http.StatusServiceUnavailable)
		return
	}
	io.WriteString(w, "ok")
}

// run runs in its own goroutine.
func (p *Proxy) run() {
	p.side = "a"
	for {
		p.poll()
		time.Sleep(30 * time.Second)
	}
}

func (p *Proxy) stop() {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.cmd != nil {
		p.cmd.Process.Kill()
	}
}

// poll runs from the run loop goroutine.
func (p *Proxy) poll() {
	heads := gerritMetaMap()
	if heads == nil {
		return
	}

	sig := p.builder.Signature(heads)

	p.mu.Lock()
	changes := sig != p.cur
	curSide := p.side
	p.cur = sig
	p.mu.Unlock()

	if !changes {
		return
	}

	newSide := "b"
	if curSide == "b" {
		newSide = "a"
	}

	dir := filepath.Join(os.TempDir(), "tip", newSide)
	if err := os.MkdirAll(dir, 0755); err != nil {
		p.err = err
		return
	}
	hostport := "localhost:8081"
	if newSide == "b" {
		hostport = "localhost:8082"
	}
	cmd, err := p.builder.Init(dir, hostport, heads)
	if err != nil {
		err = fmt.Errorf("builder.Init: %v", err)
	} else {
		go func() {
			// TODO(adg,bradfitz): be smarter about dead processes
			if err := cmd.Wait(); err != nil {
				log.Printf("process in %v exited: %v", dir, err)
			}
		}()
		err = waitReady(p.builder, hostport)
		if err != nil {
			cmd.Process.Kill()
			err = fmt.Errorf("waitReady: %v", err)
		}
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if err != nil {
		log.Println(err)
		p.err = err
		return
	}

	u, err := url.Parse(fmt.Sprintf("http://%v/", hostport))
	if err != nil {
		err = fmt.Errorf("parsing hostport: %v", err)
		log.Println(err)
		p.err = err
		return
	}
	p.proxy = httputil.NewSingleHostReverseProxy(u)
	p.side = newSide
	p.hostport = hostport
	if p.cmd != nil {
		p.cmd.Process.Kill()
	}
	p.cmd = cmd
}

func waitReady(b Builder, hostport string) error {
	var err error
	deadline := time.Now().Add(startTimeout)
	for time.Now().Before(deadline) {
		if err = b.HealthCheck(hostport); err == nil {
			return nil
		}
		time.Sleep(time.Second)
	}
	return fmt.Errorf("timed out waiting for process at %v: %v", hostport, err)
}

func runErr(cmd *exec.Cmd) error {
	out, err := cmd.CombinedOutput()
	if err != nil {
		if len(out) == 0 {
			return err
		}
		return fmt.Errorf("%s\n%v", out, err)
	}
	return nil
}

func checkout(repo, hash, path string) error {
	// Clone git repo if it doesn't exist.
	if _, err := os.Stat(filepath.Join(path, ".git")); os.IsNotExist(err) {
		if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
			return fmt.Errorf("mkdir: %v", err)
		}
		if err := runErr(exec.Command("git", "clone", repo, path)); err != nil {
			return fmt.Errorf("clone: %v", err)
		}
	} else if err != nil {
		return fmt.Errorf("stat .git: %v", err)
	}

	// Pull down changes and update to hash.
	cmd := exec.Command("git", "fetch")
	cmd.Dir = path
	if err := runErr(cmd); err != nil {
		return fmt.Errorf("fetch: %v", err)
	}
	cmd = exec.Command("git", "reset", "--hard", hash)
	cmd.Dir = path
	if err := runErr(cmd); err != nil {
		return fmt.Errorf("reset: %v", err)
	}
	cmd = exec.Command("git", "clean", "-d", "-f", "-x")
	cmd.Dir = path
	if err := runErr(cmd); err != nil {
		return fmt.Errorf("clean: %v", err)
	}
	return nil
}

// gerritMetaMap returns the map from repo name (e.g. "go") to its
// latest master hash.
// The returned map is nil on any transient error.
func gerritMetaMap() map[string]string {
	res, err := http.Get(metaURL)
	if err != nil {
		return nil
	}
	defer res.Body.Close()
	defer io.Copy(ioutil.Discard, res.Body) // ensure EOF for keep-alive
	if res.StatusCode != 200 {
		return nil
	}
	var meta map[string]struct {
		Branches map[string]string
	}
	br := bufio.NewReader(res.Body)
	// For security reasons or something, this URL starts with ")]}'\n" before
	// the JSON object. So ignore that.
	// Shawn Pearce says it's guaranteed to always be just one line, ending in '\n'.
	for {
		b, err := br.ReadByte()
		if err != nil {
			return nil
		}
		if b == '\n' {
			break
		}
	}
	if err := json.NewDecoder(br).Decode(&meta); err != nil {
		log.Printf("JSON decoding error from %v: %s", metaURL, err)
		return nil
	}
	m := map[string]string{}
	for repo, v := range meta {
		if master, ok := v.Branches["master"]; ok {
			m[repo] = master
		}
	}
	return m
}

func getOK(url string) (body []byte, err error) {
	res, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	body, err = ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return nil, err
	}
	if res.StatusCode != http.StatusOK {
		return nil, errors.New(res.Status)
	}
	return body, nil
}

// httpsOnlyHandler redirects requests to "http://example.com/foo?bar"
// to "https://example.com/foo?bar"
type httpsOnlyHandler struct {
	h http.Handler
}

func (h httpsOnlyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Header.Get("X-Appengine-Https") == "off" {
		r.URL.Scheme = "https"
		r.URL.Host = r.Host
		http.Redirect(w, r, r.URL.String(), http.StatusFound)
		return
	}
	if r.Header.Get("X-Appengine-Https") == "on" {
		// Only set this header when we're actually in production.
		w.Header().Set("Strict-Transport-Security", "max-age=31536000; preload")
	}
	h.h.ServeHTTP(w, r)
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package keybase

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"runtime/trace"
	"sync"
	"time"

	"strings"

	"github.com/keybase/client/go/externals"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/service"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	"github.com/keybase/kbfs/env"
	"github.com/keybase/kbfs/fsrpc"
	"github.com/keybase/kbfs/libkbfs"
)

var kbCtx *libkb.GlobalContext
var conn net.Conn
var startOnce sync.Once
var logSendContext libkb.LogSendContext
var kbfsConfig libkbfs.Config

type ExternalDNSNSFetcher interface {
	GetServers() []byte
}

type dnsNSFetcher struct {
	externalFetcher ExternalDNSNSFetcher
}

func newDNSNSFetcher(d ExternalDNSNSFetcher) dnsNSFetcher {
	return dnsNSFetcher{
		externalFetcher: d,
	}
}

func (d dnsNSFetcher) processExternalResult(raw []byte) []string {
	return strings.Split(string(raw), ",")
}

func (d dnsNSFetcher) GetServers() []string {
	if d.externalFetcher != nil {
		return d.processExternalResult(d.externalFetcher.GetServers())
	}
	return getDNSServers()
}

var _ libkb.DNSNameServerFetcher = dnsNSFetcher{}

// InitOnce runs the Keybase services (only runs one time)
func InitOnce(homeDir string, logFile string, runModeStr string, accessGroupOverride bool,
	dnsNSFetcher ExternalDNSNSFetcher) {
	startOnce.Do(func() {
		if err := Init(homeDir, logFile, runModeStr, accessGroupOverride, dnsNSFetcher); err != nil {
			kbCtx.Log.Errorf("Init error: %s", err)
		}
	})
}

// Init runs the Keybase services
func Init(homeDir string, logFile string, runModeStr string, accessGroupOverride bool,
	externalDNSNSFetcher ExternalDNSNSFetcher) error {
	fmt.Println("Go: Initializing")
	if logFile != "" {
		fmt.Printf("Go: Using log: %s\n", logFile)
	}

	startTrace(logFile)

	dnsNSFetcher := newDNSNSFetcher(externalDNSNSFetcher)
	dnsServers := dnsNSFetcher.GetServers()
	for _, srv := range dnsServers {
		fmt.Printf("Go: DNS Server: %s\n", srv)
	}

	kbCtx = libkb.G
	kbCtx.Init()
	kbCtx.SetServices(externals.GetServices())
	usage := libkb.Usage{
		Config:    true,
		API:       true,
		KbKeyring: true,
	}
	runMode, err := libkb.StringToRunMode(runModeStr)
	if err != nil {
		return err
	}
	config := libkb.AppConfig{
		HomeDir:                     homeDir,
		LogFile:                     logFile,
		RunMode:                     runMode,
		Debug:                       true,
		LocalRPCDebug:               "",
		VDebugSetting:               "mobile", // use empty string for same logging as desktop default
		SecurityAccessGroupOverride: accessGroupOverride,
	}
	err = kbCtx.Configure(config, usage)
	if err != nil {
		return err
	}

	svc := service.NewService(kbCtx, false)
	err = svc.StartLoopbackServer()
	if err != nil {
		return err
	}
	kbCtx.SetService()
	uir := service.NewUIRouter(kbCtx)
	kbCtx.SetUIRouter(uir)
	kbCtx.SetDNSNameServerFetcher(dnsNSFetcher)
	svc.SetupCriticalSubServices()
	svc.RunBackgroundOperations(uir)

	serviceLog := config.GetLogFile()
	logs := libkb.Logs{
		Service: serviceLog,
	}

	logSendContext = libkb.LogSendContext{
		Contextified: libkb.NewContextified(kbCtx),
		Logs:         logs,
	}

	// open the connection
	err = Reset()
	if err != nil {
		return err
	}

	go func() {
		kbfsCtx := env.NewContextFromGlobalContext(kbCtx)
		kbfsParams := libkbfs.DefaultInitParams(kbfsCtx)
		// Setting this flag will enable KBFS debug logging to always be
		// true in a mobile setting. Kill this setting if too spammy.
		// (Setting to false now 2017-08-21 PC)
		kbfsParams.Debug = false
		kbfsParams.Mode = libkbfs.InitMinimalString
		kbfsConfig, _ = libkbfs.Init(
			context.Background(), kbfsCtx, kbfsParams, serviceCn{}, func() {},
			kbCtx.Log)
	}()

	return nil
}

type serviceCn struct {
	ctx *libkb.GlobalContext
}

func (s serviceCn) NewKeybaseService(config libkbfs.Config, params libkbfs.InitParams, ctx libkbfs.Context, log logger.Logger) (libkbfs.KeybaseService, error) {
	keybaseService := libkbfs.NewKeybaseDaemonRPC(
		config, ctx, log, true, nil, nil)
	keybaseService.AddProtocols([]rpc.Protocol{
		keybase1.FsProtocol(fsrpc.NewFS(config, log)),
		// TODO: add git protocol if mobile ever needs it.
	})
	return keybaseService, nil
}

func (s serviceCn) NewCrypto(config libkbfs.Config, params libkbfs.InitParams, ctx libkbfs.Context, log logger.Logger) (libkbfs.Crypto, error) {
	return libkbfs.NewCryptoClientRPC(config, ctx), nil
}

// LogSend sends a log to Keybase
func LogSend(status string, feedback string, sendLogs bool, uiLogPath string) (string, error) {
	logSendContext.Logs.Desktop = uiLogPath
	return logSendContext.LogSend(status, feedback, sendLogs, 5*1024*1024)
}

// WriteB64 sends a base64 encoded msgpack rpc payload
func WriteB64(str string) error {
	data, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return fmt.Errorf("Base64 decode error: %s; %s", err, str)
	}
	n, err := conn.Write(data)
	if err != nil {
		return fmt.Errorf("Write error: %s", err)
	}
	if n != len(data) {
		return errors.New("Did not write all the data")
	}
	return nil
}

const targetBufferSize = 50 * 1024

// bufferSize must be divisible by 3 to ensure that we don't split
// our b64 encode across a payload boundary if we go over our buffer
// size.
const bufferSize = targetBufferSize - (targetBufferSize % 3)

// buffer for the conn.Read
var buffer = make([]byte, bufferSize)

// ReadB64 is a blocking read for base64 encoded msgpack rpc data.
// It is called serially by the mobile run loops.
func ReadB64() (string, error) {
	n, err := conn.Read(buffer)
	if n > 0 && err == nil {
		str := base64.StdEncoding.EncodeToString(buffer[0:n])
		return str, nil
	}

	if err != nil {
		// Attempt to fix the connection
		Reset()
		return "", fmt.Errorf("Read error: %s", err)
	}

	return "", nil
}

// Reset resets the socket connection
func Reset() error {
	if conn != nil {
		conn.Close()
	}

	var err error
	conn, err = kbCtx.LoopbackListener.Dial()
	if err != nil {
		return fmt.Errorf("Socket error: %s", err)
	}
	return nil
}

// Version returns semantic version string
func Version() string {
	return libkb.VersionString()
}

func startTrace(logFile string) {
	if os.Getenv("KEYBASE_TRACE_MOBILE") != "1" {
		return
	}

	tname := filepath.Join(filepath.Dir(logFile), "svctrace.out")
	f, err := os.Create(tname)
	if err != nil {
		fmt.Printf("error creating %s\n", tname)
		return
	}
	fmt.Printf("Go: starting trace %s\n", tname)
	trace.Start(f)
	go func() {
		fmt.Printf("Go: sleeping 30s for trace\n")
		time.Sleep(30 * time.Second)
		fmt.Printf("Go: stopping trace %s\n", tname)
		trace.Stop()
		time.Sleep(5 * time.Second)
		fmt.Printf("Go: trace stopped\n")
	}()
}
<commit_msg>gobind: enable git RPC<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package keybase

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"runtime/trace"
	"sync"
	"time"

	"strings"

	"github.com/keybase/client/go/externals"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/service"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	"github.com/keybase/kbfs/env"
	"github.com/keybase/kbfs/fsrpc"
	"github.com/keybase/kbfs/libgit"
	"github.com/keybase/kbfs/libkbfs"
)

var kbCtx *libkb.GlobalContext
var conn net.Conn
var startOnce sync.Once
var logSendContext libkb.LogSendContext
var kbfsConfig libkbfs.Config

type ExternalDNSNSFetcher interface {
	GetServers() []byte
}

type dnsNSFetcher struct {
	externalFetcher ExternalDNSNSFetcher
}

func newDNSNSFetcher(d ExternalDNSNSFetcher) dnsNSFetcher {
	return dnsNSFetcher{
		externalFetcher: d,
	}
}

func (d dnsNSFetcher) processExternalResult(raw []byte) []string {
	return strings.Split(string(raw), ",")
}

func (d dnsNSFetcher) GetServers() []string {
	if d.externalFetcher != nil {
		return d.processExternalResult(d.externalFetcher.GetServers())
	}
	return getDNSServers()
}

var _ libkb.DNSNameServerFetcher = dnsNSFetcher{}

// InitOnce runs the Keybase services (only runs one time)
func InitOnce(homeDir string, logFile string, runModeStr string, accessGroupOverride bool,
	dnsNSFetcher ExternalDNSNSFetcher) {
	startOnce.Do(func() {
		if err := Init(homeDir, logFile, runModeStr, accessGroupOverride, dnsNSFetcher); err != nil {
			kbCtx.Log.Errorf("Init error: %s", err)
		}
	})
}

// Init runs the Keybase services
func Init(homeDir string, logFile string, runModeStr string, accessGroupOverride bool,
	externalDNSNSFetcher ExternalDNSNSFetcher) error {
	fmt.Println("Go: Initializing")
	if logFile != "" {
		fmt.Printf("Go: Using log: %s\n", logFile)
	}

	startTrace(logFile)

	dnsNSFetcher := newDNSNSFetcher(externalDNSNSFetcher)
	dnsServers := dnsNSFetcher.GetServers()
	for _, srv := range dnsServers {
		fmt.Printf("Go: DNS Server: %s\n", srv)
	}

	kbCtx = libkb.G
	kbCtx.Init()
	kbCtx.SetServices(externals.GetServices())
	usage := libkb.Usage{
		Config:    true,
		API:       true,
		KbKeyring: true,
	}
	runMode, err := libkb.StringToRunMode(runModeStr)
	if err != nil {
		return err
	}
	config := libkb.AppConfig{
		HomeDir:                     homeDir,
		LogFile:                     logFile,
		RunMode:                     runMode,
		Debug:                       true,
		LocalRPCDebug:               "",
		VDebugSetting:               "mobile", // use empty string for same logging as desktop default
		SecurityAccessGroupOverride: accessGroupOverride,
	}
	err = kbCtx.Configure(config, usage)
	if err != nil {
		return err
	}

	svc := service.NewService(kbCtx, false)
	err = svc.StartLoopbackServer()
	if err != nil {
		return err
	}
	kbCtx.SetService()
	uir := service.NewUIRouter(kbCtx)
	kbCtx.SetUIRouter(uir)
	kbCtx.SetDNSNameServerFetcher(dnsNSFetcher)
	svc.SetupCriticalSubServices()
	svc.RunBackgroundOperations(uir)

	serviceLog := config.GetLogFile()
	logs := libkb.Logs{
		Service: serviceLog,
	}

	logSendContext = libkb.LogSendContext{
		Contextified: libkb.NewContextified(kbCtx),
		Logs:         logs,
	}

	// open the connection
	err = Reset()
	if err != nil {
		return err
	}

	go func() {
		kbfsCtx := env.NewContextFromGlobalContext(kbCtx)
		kbfsParams := libkbfs.DefaultInitParams(kbfsCtx)
		// Setting this flag will enable KBFS debug logging to always be
		// true in a mobile setting. Kill this setting if too spammy.
		// (Setting to false now 2017-08-21 PC)
		kbfsParams.Debug = false
		kbfsParams.Mode = libkbfs.InitMinimalString
		kbfsConfig, _ = libkbfs.Init(
			context.Background(), kbfsCtx, kbfsParams, serviceCn{}, func() {},
			kbCtx.Log)
	}()

	return nil
}

type serviceCn struct {
	ctx *libkb.GlobalContext
}

func (s serviceCn) NewKeybaseService(config libkbfs.Config, params libkbfs.InitParams, ctx libkbfs.Context, log logger.Logger) (libkbfs.KeybaseService, error) {
	keybaseService := libkbfs.NewKeybaseDaemonRPC(
		config, ctx, log, true, nil, nil)
	keybaseService.AddProtocols([]rpc.Protocol{
		keybase1.FsProtocol(fsrpc.NewFS(config, log)),
		keybase1.KBFSGitProtocol(libgit.NewRPCHandler(config)),
	})
	return keybaseService, nil
}

func (s serviceCn) NewCrypto(config libkbfs.Config, params libkbfs.InitParams, ctx libkbfs.Context, log logger.Logger) (libkbfs.Crypto, error) {
	return libkbfs.NewCryptoClientRPC(config, ctx), nil
}

// LogSend sends a log to Keybase
func LogSend(status string, feedback string, sendLogs bool, uiLogPath string) (string, error) {
	logSendContext.Logs.Desktop = uiLogPath
	return logSendContext.LogSend(status, feedback, sendLogs, 5*1024*1024)
}

// WriteB64 sends a base64 encoded msgpack rpc payload
func WriteB64(str string) error {
	data, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return fmt.Errorf("Base64 decode error: %s; %s", err, str)
	}
	n, err := conn.Write(data)
	if err != nil {
		return fmt.Errorf("Write error: %s", err)
	}
	if n != len(data) {
		return errors.New("Did not write all the data")
	}
	return nil
}

const targetBufferSize = 50 * 1024

// bufferSize must be divisible by 3 to ensure that we don't split
// our b64 encode across a payload boundary if we go over our buffer
// size.
const bufferSize = targetBufferSize - (targetBufferSize % 3)

// buffer for the conn.Read
var buffer = make([]byte, bufferSize)

// ReadB64 is a blocking read for base64 encoded msgpack rpc data.
// It is called serially by the mobile run loops.
func ReadB64() (string, error) {
	n, err := conn.Read(buffer)
	if n > 0 && err == nil {
		str := base64.StdEncoding.EncodeToString(buffer[0:n])
		return str, nil
	}

	if err != nil {
		// Attempt to fix the connection
		Reset()
		return "", fmt.Errorf("Read error: %s", err)
	}

	return "", nil
}

// Reset resets the socket connection
func Reset() error {
	if conn != nil {
		conn.Close()
	}

	var err error
	conn, err = kbCtx.LoopbackListener.Dial()
	if err != nil {
		return fmt.Errorf("Socket error: %s", err)
	}
	return nil
}

// Version returns semantic version string
func Version() string {
	return libkb.VersionString()
}

func startTrace(logFile string) {
	if os.Getenv("KEYBASE_TRACE_MOBILE") != "1" {
		return
	}

	tname := filepath.Join(filepath.Dir(logFile), "svctrace.out")
	f, err := os.Create(tname)
	if err != nil {
		fmt.Printf("error creating %s\n", tname)
		return
	}
	fmt.Printf("Go: starting trace %s\n", tname)
	trace.Start(f)
	go func() {
		fmt.Printf("Go: sleeping 30s for trace\n")
		time.Sleep(30 * time.Second)
		fmt.Printf("Go: stopping trace %s\n", tname)
		trace.Stop()
		time.Sleep(5 * time.Second)
		fmt.Printf("Go: trace stopped\n")
	}()
}
<|endoftext|>


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

package main

import (
	"time"

	log "github.com/inconshreveable/log15"

	"github.com/scionproto/scion/go/cert_srv/conf"
	"github.com/scionproto/scion/go/lib/addr"
	"github.com/scionproto/scion/go/lib/common"
	"github.com/scionproto/scion/go/lib/crypto/trc"
	"github.com/scionproto/scion/go/lib/ctrl"
	"github.com/scionproto/scion/go/lib/ctrl/cert_mgmt"
	"github.com/scionproto/scion/go/lib/snet"
)

var (
	// trcReqCache is an expiring cache for pending requests of TRCs.
	trcReqCache = NewReqCache(30*time.Second, 10*time.Minute, 1*time.Second)
)

type TRCHandler struct {
	conn *snet.Conn
	ia   addr.IA
}

func NewTRCHandler(conn *snet.Conn, ia addr.IA) *TRCHandler {
	return &TRCHandler{conn: conn, ia: ia}
}

// HandleReq handles TRC requests. If the TRC is not already cached and the cache-only flag is set
// or the requester is from a remote AS, the request is dropped.
func (h *TRCHandler) HandleReq(a *snet.Addr, req *cert_mgmt.TRCReq, config *conf.Conf) {
	log.Info("Received TRC request", "addr", a, "req", req)
	var t *trc.TRC
	if req.Version == cert_mgmt.NewestVersion {
		t = config.Store.GetNewestTRC(req.ISD)
	} else {
		t = config.Store.GetTRC(req.ISD, req.Version)
	}
	srcLocal := config.PublicAddr.IA.Eq(a.IA)
	if t != nil {
		if err := h.sendTRCRep(a, t); err != nil {
			log.Error("Unable to send TRC reply", "addr", a, "req", req, "err", err)
		}
	} else if !srcLocal || req.CacheOnly {
		log.Info("Dropping TRC requeset", "addr", a, "req", req, "err", "TRC not found")
	} else {
		if err := h.fetchTRC(a, req); err != nil {
			log.Error("Unable to fetch TRC", "req", req, "err", err)
		}
	}
}

// sendTRCRep creates a TRC response and sends it to the requester.
func (h *TRCHandler) sendTRCRep(a *snet.Addr, t *trc.TRC) error {
	raw, err := t.Compress()
	if err != nil {
		return err
	}
	cpld, err := ctrl.NewCertMgmtPld(&cert_mgmt.TRC{RawTRC: raw}, nil, nil)
	if err != nil {
		return err
	}
	log.Debug("Send TRC reply", "trc", t, "addr", a)
	return SendPayload(h.conn, cpld, a)
}

// fetchTRC fetches a TRC from the remote AS.
func (h *TRCHandler) fetchTRC(a *snet.Addr, req *cert_mgmt.TRCReq) error {
	key := trc.NewKey(req.ISD, req.Version).String()
	sendReq := trcReqCache.Put(key, a)
	if sendReq { // rate limit
		return h.sendTRCReq(req)
	}
	log.Info("Ignoring TRC request (same request already pending)", "addr", a, "req", req)
	return nil
}

// sendTRCReq sends a TRC request to the specified remote AS.
func (h *TRCHandler) sendTRCReq(req *cert_mgmt.TRCReq) error {
	cpld, err := ctrl.NewCertMgmtPld(req, nil, nil)
	if err != nil {
		return err
	}
	pathSet := snet.DefNetwork.PathResolver().Query(h.ia, req.IA())
	path := pathSet.GetAppPath("")
	if path == nil {
		return common.NewBasicError("Unable to find core AS", nil)
	}
	a := &snet.Addr{IA: path.Entry.Path.DstIA(), Host: addr.SvcCS}
	log.Debug("Send TRC request", "req", req, "addr", a)
	return SendPayload(h.conn, cpld, a)
}

// HandleRep handles TRC replies. Pending requests are answered and removed.
func (h *TRCHandler) HandleRep(a *snet.Addr, rep *cert_mgmt.TRC, config *conf.Conf) {
	log.Info("Received TRC reply", "addr", a, "rep", rep)
	t, err := rep.TRC()
	if err != nil {
		log.Error("Unable to parse TRC reply", "err", err)
		return
	}
	if err = config.Store.AddTRC(t, true); err != nil {
		log.Error("Unable to store TRC", "key", t.Key(), "err", err)
		return
	}
	key := t.Key()
	reqVer := chainReqCache.Pop(key.String())
	key.Ver = cert_mgmt.NewestVersion
	reqNew := chainReqCache.Pop(key.String())
	key.Ver = t.Version
	if reqVer == nil && reqNew == nil { // No pending requests
		return
	}
	cpld, err := ctrl.NewCertMgmtPld(rep, nil, nil)
	if err != nil {
		log.Error("Unable to create TRC reply", "key", t.Key(), "err", err)
		return
	}
	h.answerReqs(reqVer, cpld, key)
	h.answerReqs(reqNew, cpld, key)
}

// answerReqs responds to pending requests.
func (h *TRCHandler) answerReqs(reqs *AddrSet, cpld *ctrl.Pld, key *trc.Key) {
	if reqs == nil {
		return
	}
	for _, dst := range reqs.Addrs {
		if err := SendPayload(h.conn, cpld, dst); err != nil {
			log.Error("Unable to write TRC reply", "key", key, "err", err)
			continue
		}
	}
}
<commit_msg>Fixed misnamed TRC request caches (#1540)<commit_after>// Copyright 2017 ETH Zurich
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"time"

	log "github.com/inconshreveable/log15"

	"github.com/scionproto/scion/go/cert_srv/conf"
	"github.com/scionproto/scion/go/lib/addr"
	"github.com/scionproto/scion/go/lib/common"
	"github.com/scionproto/scion/go/lib/crypto/trc"
	"github.com/scionproto/scion/go/lib/ctrl"
	"github.com/scionproto/scion/go/lib/ctrl/cert_mgmt"
	"github.com/scionproto/scion/go/lib/snet"
)

var (
	// trcReqCache is an expiring cache for pending requests of TRCs.
	trcReqCache = NewReqCache(30*time.Second, 10*time.Minute, 1*time.Second)
)

type TRCHandler struct {
	conn *snet.Conn
	ia   addr.IA
}

func NewTRCHandler(conn *snet.Conn, ia addr.IA) *TRCHandler {
	return &TRCHandler{conn: conn, ia: ia}
}

// HandleReq handles TRC requests. If the TRC is not already cached and the cache-only flag is set
// or the requester is from a remote AS, the request is dropped.
func (h *TRCHandler) HandleReq(a *snet.Addr, req *cert_mgmt.TRCReq, config *conf.Conf) {
	log.Info("Received TRC request", "addr", a, "req", req)
	var t *trc.TRC
	if req.Version == cert_mgmt.NewestVersion {
		t = config.Store.GetNewestTRC(req.ISD)
	} else {
		t = config.Store.GetTRC(req.ISD, req.Version)
	}
	srcLocal := config.PublicAddr.IA.Eq(a.IA)
	if t != nil {
		if err := h.sendTRCRep(a, t); err != nil {
			log.Error("Unable to send TRC reply", "addr", a, "req", req, "err", err)
		}
	} else if !srcLocal || req.CacheOnly {
		log.Info("Dropping TRC requeset", "addr", a, "req", req, "err", "TRC not found")
	} else {
		if err := h.fetchTRC(a, req); err != nil {
			log.Error("Unable to fetch TRC", "req", req, "err", err)
		}
	}
}

// sendTRCRep creates a TRC response and sends it to the requester.
func (h *TRCHandler) sendTRCRep(a *snet.Addr, t *trc.TRC) error {
	raw, err := t.Compress()
	if err != nil {
		return err
	}
	cpld, err := ctrl.NewCertMgmtPld(&cert_mgmt.TRC{RawTRC: raw}, nil, nil)
	if err != nil {
		return err
	}
	log.Debug("Send TRC reply", "trc", t, "addr", a)
	return SendPayload(h.conn, cpld, a)
}

// fetchTRC fetches a TRC from the remote AS.
func (h *TRCHandler) fetchTRC(a *snet.Addr, req *cert_mgmt.TRCReq) error {
	key := trc.NewKey(req.ISD, req.Version).String()
	sendReq := trcReqCache.Put(key, a)
	if sendReq { // rate limit
		return h.sendTRCReq(req)
	}
	log.Info("Ignoring TRC request (same request already pending)", "addr", a, "req", req)
	return nil
}

// sendTRCReq sends a TRC request to the specified remote AS.
func (h *TRCHandler) sendTRCReq(req *cert_mgmt.TRCReq) error {
	cpld, err := ctrl.NewCertMgmtPld(req, nil, nil)
	if err != nil {
		return err
	}
	pathSet := snet.DefNetwork.PathResolver().Query(h.ia, req.IA())
	path := pathSet.GetAppPath("")
	if path == nil {
		return common.NewBasicError("Unable to find core AS", nil)
	}
	a := &snet.Addr{IA: path.Entry.Path.DstIA(), Host: addr.SvcCS}
	log.Debug("Send TRC request", "req", req, "addr", a)
	return SendPayload(h.conn, cpld, a)
}

// HandleRep handles TRC replies. Pending requests are answered and removed.
func (h *TRCHandler) HandleRep(a *snet.Addr, rep *cert_mgmt.TRC, config *conf.Conf) {
	log.Info("Received TRC reply", "addr", a, "rep", rep)
	t, err := rep.TRC()
	if err != nil {
		log.Error("Unable to parse TRC reply", "err", err)
		return
	}
	if err = config.Store.AddTRC(t, true); err != nil {
		log.Error("Unable to store TRC", "key", t.Key(), "err", err)
		return
	}
	key := t.Key()
	reqVer := trcReqCache.Pop(key.String())
	key.Ver = cert_mgmt.NewestVersion
	reqNew := trcReqCache.Pop(key.String())
	key.Ver = t.Version
	if reqVer == nil && reqNew == nil { // No pending requests
		return
	}
	cpld, err := ctrl.NewCertMgmtPld(rep, nil, nil)
	if err != nil {
		log.Error("Unable to create TRC reply", "key", t.Key(), "err", err)
		return
	}
	h.answerReqs(reqVer, cpld, key)
	h.answerReqs(reqNew, cpld, key)
}

// answerReqs responds to pending requests.
func (h *TRCHandler) answerReqs(reqs *AddrSet, cpld *ctrl.Pld, key *trc.Key) {
	if reqs == nil {
		return
	}
	for _, dst := range reqs.Addrs {
		if err := SendPayload(h.conn, cpld, dst); err != nil {
			log.Error("Unable to write TRC reply", "key", key, "err", err)
			continue
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"runtime"
)

const (
	VERSION = "0.41"
)

var cmdVersion = &Command{
	Run:       runVersion,
	UsageLine: "version",
	Short:     "print Weed File System version",
	Long:      `Version prints the Weed File System version`,
}

func runVersion(cmd *Command, args []string) bool {
	if len(args) != 0 {
		cmd.Usage()
	}

	fmt.Printf("version %s %s %s\n", VERSION, runtime.GOOS, runtime.GOARCH)
	return true
}
<commit_msg>v0.42<commit_after>package main

import (
	"fmt"
	"runtime"
)

const (
	VERSION = "0.42"
)

var cmdVersion = &Command{
	Run:       runVersion,
	UsageLine: "version",
	Short:     "print Weed File System version",
	Long:      `Version prints the Weed File System version`,
}

func runVersion(cmd *Command, args []string) bool {
	if len(args) != 0 {
		cmd.Usage()
	}

	fmt.Printf("version %s %s %s\n", VERSION, runtime.GOOS, runtime.GOARCH)
	return true
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.

package gofast

import "compress/gzip"
import "bytes"

func make_gzip(t *Transport, config map[string]interface{}) (uint64, tagfn, tagfn) {
	level := config["gzip.level"].(int)
	enc := func(in, out []byte) int {
		if len(in) == 0 { // empty input
			return 0
		}
		wbuf := bytes.NewBuffer(out[:0])
		writer, err := gzip.NewWriterLevel(wbuf, level)
		if err != nil {
			panic(err)
		}
		_, err = writer.Write(in)
		if err != nil {
			panic(err)
		} else if err = writer.Flush(); err != nil {
			panic(err)
		}
		return wbuf.Len()
	}
	dec := func(in, out []byte) int {
		if len(in) == 0 {
			return 0
		}
		reader, err := gzip.NewReader(bytes.NewBuffer(in))
		if err != nil {
			panic(err)
		}
		n, err := reader.Read(out)
		if err != nil {
			panic(err)
		}
		return n
	}
	return tagGzip, enc, dec
}

func init() {
	tag_factory["gzip"] = make_gzip
}
<commit_msg>bug fix.<commit_after>//  Copyright (c) 2014 Couchbase, Inc.

package gofast

import "compress/gzip"
import "bytes"

func make_gzip(t *Transport, config map[string]interface{}) (uint64, tagfn, tagfn) {
	enc := func(in, out []byte) int {
		level := config["gzip.level"].(int)
		if len(in) == 0 { // empty input
			return 0
		}
		wbuf := bytes.NewBuffer(out[:0])
		writer, err := gzip.NewWriterLevel(wbuf, level)
		if err != nil {
			panic(err)
		}
		_, err = writer.Write(in)
		if err != nil {
			panic(err)
		} else if err = writer.Flush(); err != nil {
			panic(err)
		}
		return wbuf.Len()
	}
	dec := func(in, out []byte) int {
		if len(in) == 0 {
			return 0
		}
		reader, err := gzip.NewReader(bytes.NewBuffer(in))
		if err != nil {
			panic(err)
		}
		n, err := reader.Read(out)
		if err != nil {
			panic(err)
		}
		return n
	}
	return tagGzip, enc, dec
}

func init() {
	tag_factory["gzip"] = make_gzip
}
<|endoftext|>


<commit_before>package gomemcached

import (
	"bytes"
	"io"
	"io/ioutil"
	"reflect"
	"testing"
)

func TestTapConnectFlagNameString(t *testing.T) {
	if DUMP.String() != "DUMP" {
		t.Fatalf("Expected \"DUMP\" for DUMP, got \"%v\"",
			DUMP.String())
	}

	f := TapConnectFlag(0x3)
	exp := "BACKFILL|DUMP"
	if f.String() != exp {
		t.Fatalf("Expected %q for 0x3, got %q", exp, f.String())
	}

	f = TapConnectFlag(0x212)
	exp = "DUMP|SUPPORT_ACK|0x200"
	if f.String() != exp {
		t.Fatalf("Expected %q for 0x212, got %q", exp, f.String())
	}

	f = TapConnectFlag(0xffffffff)
	f.String() // would hang if I were stupid
}

func TestTapParsers(t *testing.T) {
	tests := []struct {
		f  TapItemParser
		in []byte

		exp       interface{}
		errs      bool
		remaining int64
	}{
		// 64 bit
		{TapParseUint64, []byte{0, 0, 0, 0}, uint64(0), true, 0},
		{TapParseUint64, []byte{0, 0, 0, 0, 0, 0, 0, 0}, uint64(0), false, 0},
		{TapParseUint64, []byte{0, 0, 0, 0, 0, 0, 0, 5}, uint64(5), false, 0},
		{TapParseUint64, []byte{0, 0, 0, 0, 0, 0, 0, 5, 6, 7}, uint64(5), false, 2},
		// 16 bit
		{TapParseUint16, []byte{0}, uint16(0), true, 0},
		{TapParseUint16, []byte{0, 0}, uint16(0), false, 0},
		{TapParseUint16, []byte{0, 5}, uint16(5), false, 0},
		{TapParseUint16, []byte{0, 5, 6, 7}, uint16(5), false, 2},
		// noop
		{TapParseBool, []byte{4, 5}, true, false, 2},
		// vbucket list
		{TapParseVBList, []byte{0}, nil, true, 0},
		{TapParseVBList, []byte{0, 0}, []uint16{}, false, 0},
		{TapParseVBList, []byte{0, 0, 0}, []uint16{}, false, 1},
		{TapParseVBList, []byte{0, 0, 0, 0}, []uint16{}, false, 2},
		{TapParseVBList, []byte{0, 1, 1, 0}, []uint16{256}, false, 0},
		{TapParseVBList, []byte{0, 2, 0, 0}, nil, true, 0},
		{TapParseVBList, []byte{0, 2, 1, 0, 0, 16}, []uint16{256, 16}, false, 0},
	}

	for _, x := range tests {
		r := bytes.NewReader(x.in)
		got, err := x.f(r)

		if (err != nil) == x.errs {
			if !reflect.DeepEqual(got, x.exp) {
				t.Errorf("Expected %v, got %v for %v",
					x.exp, got, x)
			}
			n, _ := io.Copy(ioutil.Discard, r)
			if n != x.remaining {
				t.Errorf("Expected %v remaining, got %v for %v",
					x.remaining, n, x)
			}
		} else {
			t.Errorf("Error fail, got %v on %v", err, x)
		}
	}
}
<commit_msg>Added tests for tap parsing.<commit_after>package gomemcached

import (
	"bytes"
	"encoding/binary"
	"io"
	"io/ioutil"
	"reflect"
	"testing"
)

func TestTapConnectFlagNameString(t *testing.T) {
	if DUMP.String() != "DUMP" {
		t.Fatalf("Expected \"DUMP\" for DUMP, got \"%v\"",
			DUMP.String())
	}

	f := TapConnectFlag(0x3)
	exp := "BACKFILL|DUMP"
	if f.String() != exp {
		t.Fatalf("Expected %q for 0x3, got %q", exp, f.String())
	}

	f = TapConnectFlag(0x212)
	exp = "DUMP|SUPPORT_ACK|0x200"
	if f.String() != exp {
		t.Fatalf("Expected %q for 0x212, got %q", exp, f.String())
	}

	f = TapConnectFlag(0xffffffff)
	f.String() // would hang if I were stupid
}

func TestTapParsers(t *testing.T) {
	tests := []struct {
		f  TapItemParser
		in []byte

		exp       interface{}
		errs      bool
		remaining int64
	}{
		// 64 bit
		{TapParseUint64, []byte{0, 0, 0, 0}, uint64(0), true, 0},
		{TapParseUint64, []byte{0, 0, 0, 0, 0, 0, 0, 0}, uint64(0), false, 0},
		{TapParseUint64, []byte{0, 0, 0, 0, 0, 0, 0, 5}, uint64(5), false, 0},
		{TapParseUint64, []byte{0, 0, 0, 0, 0, 0, 0, 5, 6, 7}, uint64(5), false, 2},
		// 16 bit
		{TapParseUint16, []byte{0}, uint16(0), true, 0},
		{TapParseUint16, []byte{0, 0}, uint16(0), false, 0},
		{TapParseUint16, []byte{0, 5}, uint16(5), false, 0},
		{TapParseUint16, []byte{0, 5, 6, 7}, uint16(5), false, 2},
		// noop
		{TapParseBool, []byte{4, 5}, true, false, 2},
		// vbucket list
		{TapParseVBList, []byte{0}, nil, true, 0},
		{TapParseVBList, []byte{0, 0}, []uint16{}, false, 0},
		{TapParseVBList, []byte{0, 0, 0}, []uint16{}, false, 1},
		{TapParseVBList, []byte{0, 0, 0, 0}, []uint16{}, false, 2},
		{TapParseVBList, []byte{0, 1, 1, 0}, []uint16{256}, false, 0},
		{TapParseVBList, []byte{0, 2, 0, 0}, nil, true, 0},
		{TapParseVBList, []byte{0, 2, 1, 0, 0, 16}, []uint16{256, 16}, false, 0},
	}

	for _, x := range tests {
		r := bytes.NewReader(x.in)
		got, err := x.f(r)

		if (err != nil) == x.errs {
			if !reflect.DeepEqual(got, x.exp) {
				t.Errorf("Expected %v, got %v for %v",
					x.exp, got, x)
			}
			n, _ := io.Copy(ioutil.Discard, r)
			if n != x.remaining {
				t.Errorf("Expected %v remaining, got %v for %v",
					x.remaining, n, x)
			}
		} else {
			t.Errorf("Error fail, got %v on %v", err, x)
		}
	}
}

func TestParseTapCommandsEmpty(t *testing.T) {
	r := MCRequest{}
	c, err := r.ParseTapCommands()
	if err == nil {
		t.Fatalf("Expected error parsing empty tap conn, got: %v", c)
	}
}

func TestParseTapCommands(t *testing.T) {
	extras := make([]byte, 4)
	binary.BigEndian.PutUint32(extras, uint32(BACKFILL|DUMP|LIST_VBUCKETS))

	// Add our backfill thing.
	ourbf := uint64(824859588116)
	body := make([]byte, 8)
	binary.BigEndian.PutUint64(body, ourbf)
	// And a list of vbuckets
	body = append(body, 0, 3)
	body = append(body, 0, 1)
	body = append(body, 0, 2)
	body = append(body, 0, 4)
	// And an extra byte because
	body = append(body, 13)

	req := MCRequest{Key: []byte("hello"), Extras: extras, Body: body}
	c, err := req.ParseTapCommands()

	if err != nil {
		t.Fatalf("Error parsing tap commands: %v", err)
	}

	if c.Name != "hello" {
		t.Errorf("Where's our name? %v", c.Name)
	}

	// We added three things:
	if len(c.Flags) != 3 {
		t.Errorf("Expected three flags, got %v", c.Flags)
	}
	// And we've got a leftover
	if !reflect.DeepEqual(c.RemainingBody, []byte{13}) {
		t.Errorf("Didn't get our expected leftovers: %v", c.RemainingBody)
	}

	// Check the flags
	if !(c.Flags[DUMP]).(bool) {
		t.Errorf("Expected dump to be set. Wasn't.")
	}
	if c.Flags[BACKFILL].(uint64) != ourbf {
		t.Errorf("Expected bf to be %v, was %v", ourbf, c.Flags[BACKFILL])
	}
	if !reflect.DeepEqual(c.Flags[LIST_VBUCKETS], []uint16{1, 2, 4}) {
		t.Errorf("Didn't get our expected vbucket list: %v",
			c.Flags[LIST_VBUCKETS])
	}
}
<|endoftext|>


<commit_before>package gog

import (
	"bytes"
	"flag"
	"fmt"
	"os/exec"
	"strings"
	"testing"
	"time"

	"golang.org/x/tools/go/loader"
)

var stdlibPath = flag.String("test.stdlib-pkg", "", "in TestStdlib, only graph this package (import path)")

// adapted from go/types stdlib_test.go

var (
	pkgCount         int // number of packages processed
	defCount         int
	refCount         int
	unresolvedIdents int
	start            time.Time
)

func TestStdlib(t *testing.T) {
	if testing.Short() {
		t.Skip()
	}

	if data, err := exec.Command("go", "list", "std").Output(); err == nil {
		lines := bytes.Split(data, []byte{'\n'})
		start = time.Now()
		for _, line := range lines {
			if path := string(line); path != "" && !strings.HasPrefix(path, "cmd/") {
				if *stdlibPath == "" || path == *stdlibPath {
					testPkg(t, path)
				}
			}
		}
	} else {
		t.Fatal(err)
	}

	if testing.Verbose() {
		fmt.Println(pkgCount, "packages graphed in", time.Since(start))
		fmt.Printf("totals: %d defs, %d refs\n", defCount, refCount)
		if unresolvedIdents > 0 {
			t.Logf("unresolved idents: %d", unresolvedIdents)
		}
	}
}

func testPkg(t *testing.T, path string) {
	if path == "unsafe" {
		return
	}
	conf := Default
	conf.Import(path)
	prog, err := conf.Load()
	if err != nil {
		t.Fatal(path, err)
	}
	g := New(prog)

	start := time.Now()
	err = g.GraphAll()
	if err != nil {
		t.Fatal(err)
	}
	if testing.Verbose() {
		fmt.Printf("graphed %-22s\t% 4d msec   [% 6d defs, % 6d refs]\n", path, time.Since(start)/time.Millisecond, len(g.Defs), len(g.Refs))
	}
	pkgCount++
	defCount += len(g.Defs)
	refCount += len(g.Refs)

	checkAllIdents(t, g, prog)
	checkUnique(t, g, prog)
}

func checkUnique(t *testing.T, g *Grapher, prog *loader.Program) {
	defs := make(map[defPath]*Def, len(g.Defs))
	for _, s := range g.Defs {
		key := s.DefKey.defPath()
		if x, present := defs[key]; present {
			t.Errorf("def %+v %s:%d-%d already defined at %s:%d-%d", key, s.File, s.IdentSpan[0], s.IdentSpan[1], x.File, x.IdentSpan[0], x.IdentSpan[1])
		} else {
			defs[key] = s
		}
	}
}
<commit_msg>gog: only run stdlib test (which takes a long time) if TEST_STDLIB env var is truthy<commit_after>package gog

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"testing"
	"time"

	"golang.org/x/tools/go/loader"
)

var stdlibPath = flag.String("test.stdlib-pkg", "", "in TestStdlib, only graph this package (import path)")

// adapted from go/types stdlib_test.go

var (
	pkgCount         int // number of packages processed
	defCount         int
	refCount         int
	unresolvedIdents int
	start            time.Time
)

func TestStdlib(t *testing.T) {
	if v, _ := strconv.ParseBool(os.Getenv("TEST_STDLIB")); !v {
		t.Skip()
	}

	if data, err := exec.Command("go", "list", "std").Output(); err == nil {
		lines := bytes.Split(data, []byte{'\n'})
		start = time.Now()
		for _, line := range lines {
			if path := string(line); path != "" && !strings.HasPrefix(path, "cmd/") {
				if *stdlibPath == "" || path == *stdlibPath {
					testPkg(t, path)
				}
			}
		}
	} else {
		t.Fatal(err)
	}

	if testing.Verbose() {
		fmt.Println(pkgCount, "packages graphed in", time.Since(start))
		fmt.Printf("totals: %d defs, %d refs\n", defCount, refCount)
		if unresolvedIdents > 0 {
			t.Logf("unresolved idents: %d", unresolvedIdents)
		}
	}
}

func testPkg(t *testing.T, path string) {
	if path == "unsafe" {
		return
	}
	conf := Default
	conf.Import(path)
	prog, err := conf.Load()
	if err != nil {
		t.Fatal(path, err)
	}
	g := New(prog)

	start := time.Now()
	err = g.GraphAll()
	if err != nil {
		t.Fatal(err)
	}
	if testing.Verbose() {
		fmt.Printf("graphed %-22s\t% 4d msec   [% 6d defs, % 6d refs]\n", path, time.Since(start)/time.Millisecond, len(g.Defs), len(g.Refs))
	}
	pkgCount++
	defCount += len(g.Defs)
	refCount += len(g.Refs)

	checkAllIdents(t, g, prog)
	checkUnique(t, g, prog)
}

func checkUnique(t *testing.T, g *Grapher, prog *loader.Program) {
	defs := make(map[defPath]*Def, len(g.Defs))
	for _, s := range g.Defs {
		key := s.DefKey.defPath()
		if x, present := defs[key]; present {
			t.Errorf("def %+v %s:%d-%d already defined at %s:%d-%d", key, s.File, s.IdentSpan[0], s.IdentSpan[1], x.File, x.IdentSpan[0], x.IdentSpan[1])
		} else {
			defs[key] = s
		}
	}
}
<|endoftext|>


<commit_before>package model

import (
	"appengine"
	"appengine/datastore"
	"crypto/md5"
	goerrors "errors"
	"fmt"
	"github.com/knightso/base/errors"
	"github.com/qedus/nds"
	"io"
	"reflect"
	"time"
)

var OptimisticLockError = goerrors.New("Optimistic lock failure.")
var UniqueConstraintError = goerrors.New("Already exists.")

var DefaultCache = false
var CacheKinds = make(map[string]bool)

func needCache(string kind) bool {
	if needCache(kind) {
		return needCache
	} else {
		return DefaultCache
	}
}

func getGetFunc(string kind) func(appengine.Context, *datastore.Key, interface{}) error {
	if needCache(kind) {
		return nds.Get
	} else {
		return datastore.Get
	}
}

func getGetMultiFunc(string kind) func(appengine.Context, []*datastore.Key, interface{}) error {
	if needCache(kind) {
		return nds.GetMulti
	} else {
		return datastore.GetMulti
	}
}

func getPutFunc(string kind) func(appengine.Context, *datastore.Key, interface{}) (*datastore.Key, error) {
	if needCache(kind) {
		return nds.Put
	} else {
		return datastore.Put
	}
}

func getPutMultiFunc(string kind) func(appengine.Context, []*datastore.Key, interface{}) ([]*datastore.Key, error) {
	if needCache(kind) {
		return nds.PutMulti
	} else {
		return datastore.PutMulti
	}
}

func getDeleteFunc(string kind) func(appengine.Context, *datastore.Key) error {
	if needCache(kind) {
		return nds.Delete
	} else {
		return datastore.Delete
	}
}

type HasKey interface {
	GetKey() *datastore.Key
	SetKey(*datastore.Key)
}

type HasTime interface {
	SetMetaTime()
	GetCreatedAt() time.Time
	GetUpdatedAt() time.Time
}

type HasVersion interface {
	GetVersion() int
	IncrementVersion()
}

type Meta struct {
	Key       *datastore.Key `datastore:"-" json:"-"`
	Version   int            `datastore:",noindex" json:"version"`
	Deleted   bool           `json:"deleted"`
	CreatedAt time.Time      `json:"createdAt"`
	UpdatedAt time.Time      `json:"updatedAt"`
}

func (m *Meta) GetKey() *datastore.Key {
	return m.Key
}

func (m *Meta) SetKey(key *datastore.Key) {
	m.Key = key
}

func (m *Meta) SetMetaTime() {
	if m.CreatedAt == (time.Time{}) {
		m.CreatedAt = time.Now()
	}
	m.UpdatedAt = time.Now()
}

func (m *Meta) GetCreatedAt() time.Time {
	return m.CreatedAt
}

func (m *Meta) GetUpdatedAt() time.Time {
	return m.UpdatedAt
}

func (m *Meta) GetVersion() int {
	return m.Version
}

func (m *Meta) IncrementVersion() {
	m.Version++
}

func Put(c appengine.Context, e HasKey) error {
	if ht, ok := e.(HasTime); ok {
		ht.SetMetaTime()
	}
	if hv, ok := e.(HasVersion); ok {
		hv.IncrementVersion()
	}

	key, err := nds.Put(c, e.GetKey(), e)
	if err != nil {
		return errors.WrapOr(err)
	}

	e.SetKey(key)

	return nil
}

func Get(c appengine.Context, key *datastore.Key, dst interface{}) error {

	if err := nds.Get(c, key, dst); err != nil {
		return errors.WrapOr(err)
	}

	if hk, ok := dst.(HasKey); ok {
		hk.SetKey(key)
	}

	return nil
}

func GetWithVersion(c appengine.Context, key *datastore.Key, version int, dst interface{}) error {

	if err := nds.Get(c, key, dst); err != nil {
		return errors.WrapOr(err)
	}

	if hv, ok := dst.(HasVersion); ok {
		if hv.GetVersion() != version {
			return errors.WrapOr(OptimisticLockError)
		}
	}

	if hk, ok := dst.(HasKey); ok {
		hk.SetKey(key)
	}

	return nil
}

// keys only query is not supported(use datastore package)
// TODO:check keysonly
func ExecuteQuery(c appengine.Context, q *datastore.Query, dst interface{}) error {

	keys, err := q.GetAll(c, dst)
	if err != nil {
		return errors.WrapOr(err)
	}

	return forEach(dst, func(index int, elem interface{}) error {
		if hk, ok := elem.(HasKey); ok {
			hk.SetKey(keys[index])
		}
		return nil
	})
}

func GetMulti(c appengine.Context, keys []*datastore.Key, dst interface{}) error {

	err := nds.GetMulti(c, keys, dst)
	if err != nil {
		if _, ok := err.(appengine.MultiError); !ok {
			return errors.WrapOr(err)
		}
	}

	forEach(dst, func(index int, elem interface{}) error {
		if hk, ok := elem.(HasKey); ok {
			hk.SetKey(keys[index])
		}
		return nil
	})

	if err != nil {
		return errors.WrapOr(err)
	}

	return nil
}

func GetMultiWithVersion(c appengine.Context, keys []*datastore.Key, versions []int, dst interface{}) error {

	err := nds.GetMulti(c, keys, dst)
	if err != nil {
		if _, ok := err.(appengine.MultiError); !ok {
			return errors.WrapOr(err)
		}
	}

	return forEach(dst, func(index int, elem interface{}) error {
		if hk, ok := elem.(HasKey); ok {
			hk.SetKey(keys[index])
		}
		if hv, ok := elem.(HasVersion); ok {
			if hv.GetVersion() != versions[index] {
				return errors.WrapOr(OptimisticLockError)
			}
		}
		return nil
	})

	if err != nil {
		return errors.WrapOr(err)
	}

	return nil
}

// TODO: consider interface (how do you treat key?)
func PutMulti(c appengine.Context, keys []*datastore.Key, dst interface{}) error {

	forEach(dst, func(index int, elem interface{}) error {
		if ht, ok := elem.(HasTime); ok {
			ht.SetMetaTime()
		}
		if hv, ok := elem.(HasVersion); ok {
			hv.IncrementVersion()
		}
		return nil
	})

	keysput, err := nds.PutMulti(c, keys, dst)
	if err != nil {
		if _, ok := err.(appengine.MultiError); !ok {
			return errors.WrapOr(err)
		}
	}

	forEach(dst, func(index int, elem interface{}) error {
		if hk, ok := elem.(HasKey); ok {
			hk.SetKey(keysput[index])
		}
		return nil
	})

	if err != nil {
		return errors.WrapOr(err)
	}

	return nil
}

func forEach(list interface{}, f func(index int, elem interface{}) error) error {

	v := reflect.ValueOf(list)
	v = reflect.Indirect(v)
	elemKind := v.Type().Elem().Kind()

	for i := 0; i < v.Len(); i++ {
		elem := v.Index(i)

		if elemKind == reflect.Struct {
			elem = elem.Addr()
		} else if elem.IsNil() {
			continue
		}

		ei := elem.Interface()

		if err := f(i, ei); err != nil {
			return errors.WrapOr(err)
		}
	}
	return nil
}

type GenericIndex struct {
	Key *datastore.Key
}

const INDEX_PREFIX = "Index"

func FindIndexedKey(c appengine.Context, kind, id string) (*datastore.Key, error) {
	idxKey := datastore.NewKey(c, kind+INDEX_PREFIX, id, 0, nil)

	var index GenericIndex
	if err := nds.Get(c, idxKey, &index); err != nil {
		return nil, err
	}

	return index.Key, nil
}

// call in Tx
func PutKeyToIndex(c appengine.Context, key *datastore.Key, id, oldId string) error {
	if id == oldId {
		// do nothing
		return nil
	}

	idxKey := datastore.NewKey(c, key.Kind()+INDEX_PREFIX, id, 0, nil)

	var index GenericIndex
	err := nds.Get(c, idxKey, &index)
	if err == nil {
		return errors.WrapOr(UniqueConstraintError)
	} else if err != datastore.ErrNoSuchEntity {
		return errors.WrapOr(err)
	}

	idx := GenericIndex{Key: key}
	if _, err := nds.Put(c, idxKey, &idx); err != nil {
		return errors.WrapOr(err)
	}

	if oldId != "" {
		oldIdxKey := datastore.NewKey(c, key.Kind()+INDEX_PREFIX, oldId, 0, nil)
		var oldIndex GenericIndex
		if err := nds.Get(c, oldIdxKey, &oldIndex); err != nil {
			return errors.WrapOr(err)
		} else {
			//remove old index.
			if err := nds.Delete(c, oldIdxKey); err != nil {
				return errors.WrapOr(err)
			}
		}
		if !oldIndex.Key.Equal(key) {
			return errors.New("data conflict error")
		}
	}

	return nil
}

func RemoveKeyFromIndex(c appengine.Context, key *datastore.Key, id string) error {
	idxKey := datastore.NewKey(c, key.Kind()+INDEX_PREFIX, id, 0, nil)
	if err := nds.Delete(c, idxKey); err != nil {
		return errors.WrapOr(err)
	}
	return nil
}

// to avoid tab split
func AddHashPrefix(s string) string {
	h := md5.New()
	io.WriteString(h, s)
	return fmt.Sprintf("%x-%s", h.Sum(nil)[:3], s)
}

/* under construction
func FilterStartsWith(q *datastore.Query, filterStr string, value string) *datastore.Query {
	return q.Filter(filterStr + " >=", value).Filter(filterStr + " <", value + "\ufffd")
}
*/

<commit_msg>fix: added cache flag.<commit_after>package model

import (
	"appengine"
	"appengine/datastore"
	"crypto/md5"
	goerrors "errors"
	"fmt"
	"github.com/knightso/base/errors"
	"github.com/qedus/nds"
	"io"
	"reflect"
	"time"
)

var OptimisticLockError = goerrors.New("Optimistic lock failure.")
var UniqueConstraintError = goerrors.New("Already exists.")

var DefaultCache = false
var CacheKinds = make(map[string]bool)

type HasKey interface {
	GetKey() *datastore.Key
	SetKey(*datastore.Key)
}

type HasTime interface {
	SetMetaTime()
	GetCreatedAt() time.Time
	GetUpdatedAt() time.Time
}

type HasVersion interface {
	GetVersion() int
	IncrementVersion()
}

type Meta struct {
	Key       *datastore.Key `datastore:"-" json:"-"`
	Version   int            `datastore:",noindex" json:"version"`
	Deleted   bool           `json:"deleted"`
	CreatedAt time.Time      `json:"createdAt"`
	UpdatedAt time.Time      `json:"updatedAt"`
}

func (m *Meta) GetKey() *datastore.Key {
	return m.Key
}

func (m *Meta) SetKey(key *datastore.Key) {
	m.Key = key
}

func (m *Meta) SetMetaTime() {
	if m.CreatedAt == (time.Time{}) {
		m.CreatedAt = time.Now()
	}
	m.UpdatedAt = time.Now()
}

func (m *Meta) GetCreatedAt() time.Time {
	return m.CreatedAt
}

func (m *Meta) GetUpdatedAt() time.Time {
	return m.UpdatedAt
}

func (m *Meta) GetVersion() int {
	return m.Version
}

func (m *Meta) IncrementVersion() {
	m.Version++
}

func Put(c appengine.Context, e HasKey) error {
	if ht, ok := e.(HasTime); ok {
		ht.SetMetaTime()
	}
	if hv, ok := e.(HasVersion); ok {
		hv.IncrementVersion()
	}

	f := getPutFunc(e.GetKey().Kind())
	key, err := f(c, e.GetKey(), e)
	if err != nil {
		return errors.WrapOr(err)
	}

	e.SetKey(key)

	return nil
}

func Get(c appengine.Context, key *datastore.Key, dst interface{}) error {

	f := getGetFunc(key.Kind())

	if err := f(c, key, dst); err != nil {
		return errors.WrapOr(err)
	}

	if hk, ok := dst.(HasKey); ok {
		hk.SetKey(key)
	}

	return nil
}

func GetWithVersion(c appengine.Context, key *datastore.Key, version int, dst interface{}) error {

	f := getGetFunc(key.Kind())

	if err := f(c, key, dst); err != nil {
		return errors.WrapOr(err)
	}

	if hv, ok := dst.(HasVersion); ok {
		if hv.GetVersion() != version {
			return errors.WrapOr(OptimisticLockError)
		}
	}

	if hk, ok := dst.(HasKey); ok {
		hk.SetKey(key)
	}

	return nil
}

// keys only query is not supported(use datastore package)
// TODO:check keysonly
func ExecuteQuery(c appengine.Context, q *datastore.Query, dst interface{}) error {

	keys, err := q.GetAll(c, dst)
	if err != nil {
		return errors.WrapOr(err)
	}

	return forEach(dst, func(index int, elem interface{}) error {
		if hk, ok := elem.(HasKey); ok {
			hk.SetKey(keys[index])
		}
		return nil
	})
}

func GetMulti(c appengine.Context, keys []*datastore.Key, dst interface{}) error {

	f := datastore.GetMulti
	if len(keys) > 0 {
		f = getGetMultiFunc(keys[0].Kind())
	}

	err := f(c, keys, dst)
	if err != nil {
		if _, ok := err.(appengine.MultiError); !ok {
			return errors.WrapOr(err)
		}
	}

	forEach(dst, func(index int, elem interface{}) error {
		if hk, ok := elem.(HasKey); ok {
			hk.SetKey(keys[index])
		}
		return nil
	})

	if err != nil {
		return errors.WrapOr(err)
	}

	return nil
}

func GetMultiWithVersion(c appengine.Context, keys []*datastore.Key, versions []int, dst interface{}) error {

	f := datastore.GetMulti
	if len(keys) > 0 {
		f = getGetMultiFunc(keys[0].Kind())
	}

	err := f(c, keys, dst)
	if err != nil {
		if _, ok := err.(appengine.MultiError); !ok {
			return errors.WrapOr(err)
		}
	}

	return forEach(dst, func(index int, elem interface{}) error {
		if hk, ok := elem.(HasKey); ok {
			hk.SetKey(keys[index])
		}
		if hv, ok := elem.(HasVersion); ok {
			if hv.GetVersion() != versions[index] {
				return errors.WrapOr(OptimisticLockError)
			}
		}
		return nil
	})

	if err != nil {
		return errors.WrapOr(err)
	}

	return nil
}

// TODO: consider interface (how do you treat key?)
func PutMulti(c appengine.Context, keys []*datastore.Key, dst interface{}) error {

	f := datastore.PutMulti
	if len(keys) > 0 {
		f = getPutMultiFunc(keys[0].Kind())
	}

	forEach(dst, func(index int, elem interface{}) error {
		if ht, ok := elem.(HasTime); ok {
			ht.SetMetaTime()
		}
		if hv, ok := elem.(HasVersion); ok {
			hv.IncrementVersion()
		}
		return nil
	})

	keysput, err := f(c, keys, dst)
	if err != nil {
		if _, ok := err.(appengine.MultiError); !ok {
			return errors.WrapOr(err)
		}
	}

	forEach(dst, func(index int, elem interface{}) error {
		if hk, ok := elem.(HasKey); ok {
			hk.SetKey(keysput[index])
		}
		return nil
	})

	if err != nil {
		return errors.WrapOr(err)
	}

	return nil
}

func forEach(list interface{}, f func(index int, elem interface{}) error) error {

	v := reflect.ValueOf(list)
	v = reflect.Indirect(v)
	elemKind := v.Type().Elem().Kind()

	for i := 0; i < v.Len(); i++ {
		elem := v.Index(i)

		if elemKind == reflect.Struct {
			elem = elem.Addr()
		} else if elem.IsNil() {
			continue
		}

		ei := elem.Interface()

		if err := f(i, ei); err != nil {
			return errors.WrapOr(err)
		}
	}
	return nil
}

type GenericIndex struct {
	Key *datastore.Key
}

const INDEX_PREFIX = "Index"

func FindIndexedKey(c appengine.Context, kind, id string) (*datastore.Key, error) {

	f := getGetFunc(kind)

	idxKey := datastore.NewKey(c, kind+INDEX_PREFIX, id, 0, nil)

	var index GenericIndex
	if err := f(c, idxKey, &index); err != nil {
		return nil, err
	}

	return index.Key, nil
}

// call in Tx
func PutKeyToIndex(c appengine.Context, key *datastore.Key, id, oldId string) error {
	kind := key.Kind()
	getF := getGetFunc(kind)
	putF := getPutFunc(kind)
	deleteF := getDeleteFunc(kind)

	if id == oldId {
		// do nothing
		return nil
	}

	idxKey := datastore.NewKey(c, key.Kind()+INDEX_PREFIX, id, 0, nil)

	var index GenericIndex
	err := getF(c, idxKey, &index)
	if err == nil {
		return errors.WrapOr(UniqueConstraintError)
	} else if err != datastore.ErrNoSuchEntity {
		return errors.WrapOr(err)
	}

	idx := GenericIndex{Key: key}
	if _, err := putF(c, idxKey, &idx); err != nil {
		return errors.WrapOr(err)
	}

	if oldId != "" {
		oldIdxKey := datastore.NewKey(c, key.Kind()+INDEX_PREFIX, oldId, 0, nil)
		var oldIndex GenericIndex
		if err := getF(c, oldIdxKey, &oldIndex); err != nil {
			return errors.WrapOr(err)
		} else {
			//remove old index.
			if err := deleteF(c, oldIdxKey); err != nil {
				return errors.WrapOr(err)
			}
		}
		if !oldIndex.Key.Equal(key) {
			return errors.New("data conflict error")
		}
	}

	return nil
}

func RemoveKeyFromIndex(c appengine.Context, key *datastore.Key, id string) error {
	deleteF := getDeleteFunc(key.Kind())

	idxKey := datastore.NewKey(c, key.Kind()+INDEX_PREFIX, id, 0, nil)
	if err := deleteF(c, idxKey); err != nil {
		return errors.WrapOr(err)
	}
	return nil
}

// to avoid tab split
func AddHashPrefix(s string) string {
	h := md5.New()
	io.WriteString(h, s)
	return fmt.Sprintf("%x-%s", h.Sum(nil)[:3], s)
}

func needCache(kind string) bool {
	nc, ok := CacheKinds[kind]
	if ok {
		return nc
	} else {
		return DefaultCache
	}
}

func getGetFunc(kind string) func(appengine.Context, *datastore.Key, interface{}) error {
	if needCache(kind) {
		return nds.Get
	} else {
		return datastore.Get
	}
}

func getGetMultiFunc(kind string) func(appengine.Context, []*datastore.Key, interface{}) error {
	if needCache(kind) {
		return nds.GetMulti
	} else {
		return datastore.GetMulti
	}
}

func getPutFunc(kind string) func(appengine.Context, *datastore.Key, interface{}) (*datastore.Key, error) {
	if needCache(kind) {
		return nds.Put
	} else {
		return datastore.Put
	}
}

func getPutMultiFunc(kind string) func(appengine.Context, []*datastore.Key, interface{}) ([]*datastore.Key, error) {
	if needCache(kind) {
		return nds.PutMulti
	} else {
		return datastore.PutMulti
	}
}

func getDeleteFunc(kind string) func(appengine.Context, *datastore.Key) error {
	if needCache(kind) {
		return nds.Delete
	} else {
		return datastore.Delete
	}
}

/* under construction
func FilterStartsWith(q *datastore.Query, filterStr string, value string) *datastore.Query {
	return q.Filter(filterStr + " >=", value).Filter(filterStr + " <", value + "\ufffd")
}
*/

<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
)

var (
	telegramBot *TelegramBot
)

// TelegramBot ...
type TelegramBot struct {
	Name          string
	SelfChatID    int64
	ChannelChatID int64
	ComicPath     string
	DeleteDelay   time.Duration
	Client        *tgbotapi.BotAPI
	Queue         *bt.Pool
	Tube          string
}

// NewTelegramBot ...
func NewTelegramBot(cfg *TelegramConfig, btdAddr string) (t *TelegramBot) {
	bot, err := tgbotapi.NewBotAPI(cfg.Token)
	if err != nil {
		logger.Panicf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.DeleteDelay)
	if err != nil {
		logger.Panicf("delete delay error: %+v", err)
	}

	t = &TelegramBot{
		Name:          bot.Self.UserName,
		SelfChatID:    cfg.SelfChatID,
		ChannelChatID: cfg.ChannelChatID,
		ComicPath:     cfg.ComicPath,
		DeleteDelay:   delay,
		Client:        bot,
		Tube:          "tg",
	}
	t.Queue = &bt.Pool{
		Dial: func() (*bt.Conn, error) {
			return bt.Dial(btdAddr)
		},
		MaxIdle:     10,
		MaxActive:   100,
		IdleTimeout: 60 * time.Second,
		MaxLifetime: 180 * time.Second,
		Wait:        true,
	}
	return
}

func (t *TelegramBot) putQueue(msg []byte) {
	conn, err := t.Queue.Get()
	if err != nil {
		logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(t.Tube)
	_, err = conn.Put(msg, 1, t.DeleteDelay, time.Minute)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
}

func (t *TelegramBot) send(chat int64, msg string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, msg)
	return t.Client.Send(tgbotapi.NewMessage(chat, msg))
}

func (t *TelegramBot) delMessage() {
	for {
		conn, err := t.Queue.Get()
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(t.Tube)
		job, err := conn.Reserve()
		if err != nil {
			logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = t.Client.DeleteMessage(delMsg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		t.Queue.Release(conn, false)
	}
}

func (t *TelegramBot) tgBot() {
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := t.Client.GetUpdatesChan(u)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else if update.CallbackQuery != nil {
				logger.Infof(
					"recv:(%s)[%s]reaction:{%s}",
					update.CallbackQuery.ChatInstance,
					update.CallbackQuery.From.String(),
					update.CallbackQuery.Data,
				)
				data := strings.SplitN(update.CallbackQuery.Data, ":", 2)
				switch data[0] {
				case "comic", "pic":
					go onReaction(t, update.CallbackQuery)
				}
				continue
			} else {
				continue
			}
			if message.Chat.IsGroup() {
				logger.Infof(
					"recv:(%s)[%s]{%s}",
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				logger.Infof(
					"recv:[%s]{%s}",
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(t, message)
				case "comic":
					go onComic(t, message)
				case "pic":
					go onPic(t, message)
				default:
					logger.Infof("ignore unkown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				checkRepeat(t, message)
			}
		}
		logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func checkRepeat(t *TelegramBot, message *tgbotapi.Message) {
	key := "tg_" + getMsgTitle(message) + "_last"
	flattendMsg := strings.TrimSpace(message.Text)
	defer redisClient.LTrim(key, 0, 10)
	defer redisClient.LPush(key, flattendMsg)

	lastMsgs, err := redisClient.LRange(key, 0, 6).Result()
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	i := 0
	for _, s := range lastMsgs {
		if s == flattendMsg {
			i++
		}
	}
	if i > 1 {
		redisClient.Del(key)
		logger.Infof("repeat: %s", strconv.Quote(message.Text))
		msg := tgbotapi.NewMessage(message.Chat.ID, message.Text)
		t.Client.Send(msg)
	}
}

func onStart(t *TelegramBot, message *tgbotapi.Message) {
	msg := tgbotapi.NewMessage(message.Chat.ID, "呀呀呀")
	msg.ReplyToMessageID = message.MessageID
	t.Client.Send(msg)
}

func onComic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(t.ComicPath)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]
	number := strings.Split(strings.Split(file, "@")[1], ".")[0]
	msg := tgbotapi.NewMessage(message.Chat.ID, "🔞 https://nhentai.net/g/"+number)

	msg.ReplyMarkup = buildInlineKeyboardMarkup("comic", number)

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func onPic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(twitterBot.ImgPath, "*"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pics")
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))

	msg := tgbotapi.NewDocumentUpload(message.Chat.ID, file)
	msg.ReplyMarkup = buildInlineKeyboardMarkup("pic", filepath.Base(file))

	_, err = t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func onReaction(t *TelegramBot, callbackQuery *tgbotapi.CallbackQuery) {
	var callbackText string

	_type, _id, reaction, err := saveReaction(callbackQuery.Data, callbackQuery.From.ID)
	if err == nil {
		msg := tgbotapi.NewEditMessageReplyMarkup(
			callbackQuery.Message.Chat.ID,
			callbackQuery.Message.MessageID,
			buildInlineKeyboardMarkup(_type, _id),
		)
		_, err = t.Client.Send(msg)
	}

	if err != nil {
		logger.Errorf("%+v", err)
		callbackText = err.Error()
	} else {
		callbackText = reaction + " " + _id + "!"
	}

	callbackMsg := tgbotapi.NewCallback(callbackQuery.ID, callbackText)
	_, err = t.Client.AnswerCallbackQuery(callbackMsg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}

func buildReactionData(_type, _id, reaction string) string {
	return _type + ":" + _id + ":" + reaction
}
func buildReactionKey(_type, _id, reaction string) string {
	return "tg_reaction_" + buildReactionData(_type, _id, reaction)
}

func buildInlineKeyboardMarkup(_type, _id string) tgbotapi.InlineKeyboardMarkup {

	likeCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "like")).Result()
	dissCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Result()

	likeText := "❤"
	if likeCount > 0 {
		likeText = likeText + " " + strconv.FormatInt(likeCount, 10)
	}
	dissText := "💔"
	if dissCount > 0 {
		dissText = dissText + " " + strconv.FormatInt(dissCount, 10)
	}

	row := tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData(likeText, buildReactionData(_type, _id, "like")),
		tgbotapi.NewInlineKeyboardButtonData(dissText, buildReactionData(_type, _id, "diss")),
	)
	return tgbotapi.NewInlineKeyboardMarkup(row)
}

func saveReaction(key string, user int) (_type, _id, reaction string, err error) {
	token := strings.Split(key, ":")
	if len(token) != 3 {
		err = fmt.Errorf("react data error: %s", key)
		return
	}
	_type = token[0]
	_id = token[1]
	reaction = token[2]

	pipe := redisClient.Pipeline()
	switch reaction {
	case "like":
		likeCount := pipe.SAdd(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		dissCount := pipe.SRem(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	case "diss":
		dissCount := pipe.SAdd(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		likeCount := pipe.SRem(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	default:
		err = fmt.Errorf("react type error: %s", key)
	}
	return
}
<commit_msg>log debug instead of error<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
)

var (
	telegramBot *TelegramBot
)

// TelegramBot ...
type TelegramBot struct {
	Name          string
	SelfChatID    int64
	ChannelChatID int64
	ComicPath     string
	DeleteDelay   time.Duration
	Client        *tgbotapi.BotAPI
	Queue         *bt.Pool
	Tube          string
}

// NewTelegramBot ...
func NewTelegramBot(cfg *TelegramConfig, btdAddr string) (t *TelegramBot) {
	bot, err := tgbotapi.NewBotAPI(cfg.Token)
	if err != nil {
		logger.Panicf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.DeleteDelay)
	if err != nil {
		logger.Panicf("delete delay error: %+v", err)
	}

	t = &TelegramBot{
		Name:          bot.Self.UserName,
		SelfChatID:    cfg.SelfChatID,
		ChannelChatID: cfg.ChannelChatID,
		ComicPath:     cfg.ComicPath,
		DeleteDelay:   delay,
		Client:        bot,
		Tube:          "tg",
	}
	t.Queue = &bt.Pool{
		Dial: func() (*bt.Conn, error) {
			return bt.Dial(btdAddr)
		},
		MaxIdle:     10,
		MaxActive:   100,
		IdleTimeout: 60 * time.Second,
		MaxLifetime: 180 * time.Second,
		Wait:        true,
	}
	return
}

func (t *TelegramBot) putQueue(msg []byte) {
	conn, err := t.Queue.Get()
	if err != nil {
		logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(t.Tube)
	_, err = conn.Put(msg, 1, t.DeleteDelay, time.Minute)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
}

func (t *TelegramBot) send(chat int64, msg string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, msg)
	return t.Client.Send(tgbotapi.NewMessage(chat, msg))
}

func (t *TelegramBot) delMessage() {
	for {
		conn, err := t.Queue.Get()
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(t.Tube)
		job, err := conn.Reserve()
		if err != nil {
			logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = t.Client.DeleteMessage(delMsg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		t.Queue.Release(conn, false)
	}
}

func (t *TelegramBot) tgBot() {
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := t.Client.GetUpdatesChan(u)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else if update.CallbackQuery != nil {
				logger.Infof(
					"recv:(%s)[%s]reaction:{%s}",
					update.CallbackQuery.ChatInstance,
					update.CallbackQuery.From.String(),
					update.CallbackQuery.Data,
				)
				data := strings.SplitN(update.CallbackQuery.Data, ":", 2)
				switch data[0] {
				case "comic", "pic":
					go onReaction(t, update.CallbackQuery)
				}
				continue
			} else {
				continue
			}
			if message.Chat.IsGroup() {
				logger.Infof(
					"recv:(%s)[%s]{%s}",
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				logger.Infof(
					"recv:[%s]{%s}",
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(t, message)
				case "comic":
					go onComic(t, message)
				case "pic":
					go onPic(t, message)
				default:
					logger.Infof("ignore unkown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				checkRepeat(t, message)
			}
		}
		logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func checkRepeat(t *TelegramBot, message *tgbotapi.Message) {
	key := "tg_" + getMsgTitle(message) + "_last"
	flattendMsg := strings.TrimSpace(message.Text)
	defer redisClient.LTrim(key, 0, 10)
	defer redisClient.LPush(key, flattendMsg)

	lastMsgs, err := redisClient.LRange(key, 0, 6).Result()
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	i := 0
	for _, s := range lastMsgs {
		if s == flattendMsg {
			i++
		}
	}
	if i > 1 {
		redisClient.Del(key)
		logger.Infof("repeat: %s", strconv.Quote(message.Text))
		msg := tgbotapi.NewMessage(message.Chat.ID, message.Text)
		t.Client.Send(msg)
	}
}

func onStart(t *TelegramBot, message *tgbotapi.Message) {
	msg := tgbotapi.NewMessage(message.Chat.ID, "呀呀呀")
	msg.ReplyToMessageID = message.MessageID
	t.Client.Send(msg)
}

func onComic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(t.ComicPath)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]
	number := strings.Split(strings.Split(file, "@")[1], ".")[0]
	msg := tgbotapi.NewMessage(message.Chat.ID, "🔞 https://nhentai.net/g/"+number)

	msg.ReplyMarkup = buildInlineKeyboardMarkup("comic", number)

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func onPic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(twitterBot.ImgPath, "*"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pics")
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))

	msg := tgbotapi.NewDocumentUpload(message.Chat.ID, file)
	msg.ReplyMarkup = buildInlineKeyboardMarkup("pic", filepath.Base(file))

	_, err = t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func onReaction(t *TelegramBot, callbackQuery *tgbotapi.CallbackQuery) {
	var callbackText string

	_type, _id, reaction, err := saveReaction(callbackQuery.Data, callbackQuery.From.ID)
	if err == nil {
		msg := tgbotapi.NewEditMessageReplyMarkup(
			callbackQuery.Message.Chat.ID,
			callbackQuery.Message.MessageID,
			buildInlineKeyboardMarkup(_type, _id),
		)
		_, err = t.Client.Send(msg)
	}

	if err != nil {
		logger.Debugf("%+v", err)
		callbackText = err.Error()
	} else {
		callbackText = reaction + " " + _id + "!"
	}

	callbackMsg := tgbotapi.NewCallback(callbackQuery.ID, callbackText)
	_, err = t.Client.AnswerCallbackQuery(callbackMsg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}

func buildReactionData(_type, _id, reaction string) string {
	return _type + ":" + _id + ":" + reaction
}
func buildReactionKey(_type, _id, reaction string) string {
	return "tg_reaction_" + buildReactionData(_type, _id, reaction)
}

func buildInlineKeyboardMarkup(_type, _id string) tgbotapi.InlineKeyboardMarkup {

	likeCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "like")).Result()
	dissCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Result()

	likeText := "❤"
	if likeCount > 0 {
		likeText = likeText + " " + strconv.FormatInt(likeCount, 10)
	}
	dissText := "💔"
	if dissCount > 0 {
		dissText = dissText + " " + strconv.FormatInt(dissCount, 10)
	}

	row := tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData(likeText, buildReactionData(_type, _id, "like")),
		tgbotapi.NewInlineKeyboardButtonData(dissText, buildReactionData(_type, _id, "diss")),
	)
	return tgbotapi.NewInlineKeyboardMarkup(row)
}

func saveReaction(key string, user int) (_type, _id, reaction string, err error) {
	token := strings.Split(key, ":")
	if len(token) != 3 {
		err = fmt.Errorf("react data error: %s", key)
		return
	}
	_type = token[0]
	_id = token[1]
	reaction = token[2]

	pipe := redisClient.Pipeline()
	switch reaction {
	case "like":
		likeCount := pipe.SAdd(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		dissCount := pipe.SRem(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	case "diss":
		dissCount := pipe.SAdd(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		likeCount := pipe.SRem(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	default:
		err = fmt.Errorf("react type error: %s", key)
	}
	return
}
<|endoftext|>


<commit_before>package gateway

const VERSION = "v2.9.1"
<commit_msg>Bump version to 2.9.2 (#2700)<commit_after>package gateway

const VERSION = "v2.9.2"
<|endoftext|>


<commit_before>package sql

import (
	"fmt"
	"io"
)

// SelectStatement represents a SQL SELECT statement.
type SelectStatement struct {
	Fields    []string
	TableName string
}

type InsertStatement struct {
	MapValues map[string]interface{}
	Values    []interface{}
	TableName string
}

// Parser represents a parser.
type Parser struct {
	s   *Scanner
	buf struct {
		tok Token  // last read token
		lit string // last read literal
		n   int    // buffer size (max=1)
	}
}

// NewParser returns a new instance of Parser.
func NewParser(r io.Reader) *Parser {
	return &Parser{s: NewScanner(r)}
}

// Parse parses a SQL SELECT statement.
func (p *Parser) Parse() (interface{}, error) {

	//Find the first token
	tok, lit := p.scanIgnoreWhitespace()

	switch tok {
	case SELECT:
		p.unscan() //We unscan to get back to what we've already parsed
		return p.ParseSelect()
	case INSERT:
		p.unscan()
		return p.ParseInsert()
	default:
		return nil, fmt.Errorf("found %q, expected SELECT/INSERT/CREATE", lit)
	}

	// Return the successfully parsed statement.
	return nil, fmt.Errorf("Could not parse the statement")
}

//ParseInsert parses the given string and returns an instance of InsertStatement
func (p *Parser) ParseInsert() (InsertStatement, error) {
	stmt := &InsertStatement{}

	if tok, lit := p.scanIgnoreWhitespace(); tok != INSERT {
		return *stmt, fmt.Errorf("found %q, expected INSERT", lit)
	}

	if tok, lit := p.scanIgnoreWhitespace(); tok != INTO {
		return *stmt, fmt.Errorf("found %q, expected INTO", lit)
	}

	if tok, lit := p.scanIgnoreWhitespace(); tokenMap[tok] {
		return *stmt, fmt.Errorf("found %q, table name cannot be a keyword", lit)
	} else {
		stmt.TableName = lit
	}

	return *stmt, nil
}

//ParseSelect parses the given string and returns an instance of SelectStatement
func (p *Parser) ParseSelect() (SelectStatement, error) {
	stmt := &SelectStatement{}
	// First token should be a "SELECT" keyword.
	if tok, lit := p.scanIgnoreWhitespace(); tok != SELECT {
		return *stmt, fmt.Errorf("found %q, expected SELECT", lit)
	}

	// Next we should loop over all our comma-delimited fields.
	for {
		// Read a field.
		tok, lit := p.scanIgnoreWhitespace()
		if tok != IDENT && tok != ASTERISK {
			return *stmt, fmt.Errorf("found %q, expected field", lit)
		}
		stmt.Fields = append(stmt.Fields, lit)

		// If the next token is not a comma then break the loop.
		if tok, _ := p.scanIgnoreWhitespace(); tok != COMMA {
			p.unscan()
			break
		}
	}

	// Next we should see the "FROM" keyword.
	if tok, lit := p.scanIgnoreWhitespace(); tok != FROM {
		return *stmt, fmt.Errorf("found %q, expected FROM", lit)
	}

	// Finally we should read the table name.
	tok, lit := p.scanIgnoreWhitespace()
	if tok != IDENT {
		return *stmt, fmt.Errorf("found %q, expected table name", lit)
	}
	stmt.TableName = lit

	// Return the successfully parsed statement.
	return *stmt, nil
}

// scan returns the next token from the underlying scanner.
// If a token has been unscanned then read that instead.
func (p *Parser) scan() (tok Token, lit string) {
	// If we have a token on the buffer, then return it.
	if p.buf.n != 0 {
		p.buf.n = 0
		return p.buf.tok, p.buf.lit
	}

	// Otherwise read the next token from the scanner.
	tok, lit = p.s.Scan()

	// Save it to the buffer in case we unscan later.
	p.buf.tok, p.buf.lit = tok, lit

	return
}

// scanIgnoreWhitespace scans the next non-whitespace token.
func (p *Parser) scanIgnoreWhitespace() (tok Token, lit string) {
	tok, lit = p.scan()
	if tok == WS {
		tok, lit = p.scan()
	}
	return
}

// unscan pushes the previously read token back onto the buffer.
func (p *Parser) unscan() { p.buf.n = 1 }
<commit_msg>Use IDENT instead of the tokenMap<commit_after>package sql

import (
	"fmt"
	"io"
)

// SelectStatement represents a SQL SELECT statement.
type SelectStatement struct {
	Fields    []string
	TableName string
}

type InsertStatement struct {
	MapValues map[string]interface{}
	Values    []interface{}
	TableName string
}

// Parser represents a parser.
type Parser struct {
	s   *Scanner
	buf struct {
		tok Token  // last read token
		lit string // last read literal
		n   int    // buffer size (max=1)
	}
}

// NewParser returns a new instance of Parser.
func NewParser(r io.Reader) *Parser {
	return &Parser{s: NewScanner(r)}
}

// Parse parses a SQL SELECT statement.
func (p *Parser) Parse() (interface{}, error) {

	//Find the first token
	tok, lit := p.scanIgnoreWhitespace()

	switch tok {
	case SELECT:
		p.unscan() //We unscan to get back to what we've already parsed
		return p.ParseSelect()
	case INSERT:
		p.unscan()
		return p.ParseInsert()
	default:
		return nil, fmt.Errorf("found %q, expected SELECT/INSERT/CREATE", lit)
	}

	// Return the successfully parsed statement.
	return nil, fmt.Errorf("Could not parse the statement")
}

//ParseInsert parses the given string and returns an instance of InsertStatement
func (p *Parser) ParseInsert() (InsertStatement, error) {
	stmt := &InsertStatement{}

	if tok, lit := p.scanIgnoreWhitespace(); tok != INSERT {
		return *stmt, fmt.Errorf("found %q, expected INSERT", lit)
	}

	if tok, lit := p.scanIgnoreWhitespace(); tok != INTO {
		return *stmt, fmt.Errorf("found %q, expected INTO", lit)
	}

	if tok, lit := p.scanIgnoreWhitespace(); tok != IDENT {
		return *stmt, fmt.Errorf("found %q, table name cannot be a keyword", lit)
	} else {
		stmt.TableName = lit
	}

	return *stmt, nil
}

//ParseSelect parses the given string and returns an instance of SelectStatement
func (p *Parser) ParseSelect() (SelectStatement, error) {
	stmt := &SelectStatement{}
	// First token should be a "SELECT" keyword.
	if tok, lit := p.scanIgnoreWhitespace(); tok != SELECT {
		return *stmt, fmt.Errorf("found %q, expected SELECT", lit)
	}

	// Next we should loop over all our comma-delimited fields.
	for {
		// Read a field.
		tok, lit := p.scanIgnoreWhitespace()
		if tok != IDENT && tok != ASTERISK {
			return *stmt, fmt.Errorf("found %q, expected field", lit)
		}
		stmt.Fields = append(stmt.Fields, lit)

		// If the next token is not a comma then break the loop.
		if tok, _ := p.scanIgnoreWhitespace(); tok != COMMA {
			p.unscan()
			break
		}
	}

	// Next we should see the "FROM" keyword.
	if tok, lit := p.scanIgnoreWhitespace(); tok != FROM {
		return *stmt, fmt.Errorf("found %q, expected FROM", lit)
	}

	// Finally we should read the table name.
	tok, lit := p.scanIgnoreWhitespace()
	if tok != IDENT {
		return *stmt, fmt.Errorf("found %q, expected table name", lit)
	}
	stmt.TableName = lit

	// Return the successfully parsed statement.
	return *stmt, nil
}

// scan returns the next token from the underlying scanner.
// If a token has been unscanned then read that instead.
func (p *Parser) scan() (tok Token, lit string) {
	// If we have a token on the buffer, then return it.
	if p.buf.n != 0 {
		p.buf.n = 0
		return p.buf.tok, p.buf.lit
	}

	// Otherwise read the next token from the scanner.
	tok, lit = p.s.Scan()

	// Save it to the buffer in case we unscan later.
	p.buf.tok, p.buf.lit = tok, lit

	return
}

// scanIgnoreWhitespace scans the next non-whitespace token.
func (p *Parser) scanIgnoreWhitespace() (tok Token, lit string) {
	tok, lit = p.scan()
	if tok == WS {
		tok, lit = p.scan()
	}
	return
}

// unscan pushes the previously read token back onto the buffer.
func (p *Parser) unscan() { p.buf.n = 1 }
<|endoftext|>