<commit_before>package common

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// errors
var (
	ErrTwitchAlreadyInChannel = errors.New("already in channel")
	ErrTwitchNotInChannel     = errors.New("not in channel")
)

// TwitchChat twitch chat client
type TwitchChat struct {
	sync.RWMutex
	conn        *websocket.Conn
	dialer      websocket.Dialer
	headers     http.Header
	messages    map[string]chan *Message
	channels    []string
	writeLock   sync.Mutex
	joinHandler TwitchJoinHandler
	admins      map[string]bool
}

// TwitchJoinHandler called when joining channels
type TwitchJoinHandler func(string, chan *Message)

// NewTwitchChat new twitch chat client
func NewTwitchChat(j TwitchJoinHandler) *TwitchChat {
	c := &TwitchChat{
		dialer:      websocket.Dialer{HandshakeTimeout: SocketHandshakeTimeout},
		headers:     http.Header{"Origin": []string{GetConfig().Twitch.OriginURL}},
		messages:    make(map[string]chan *Message, 0),
		channels:    make([]string, 0),
		joinHandler: j,
		admins:      make(map[string]bool, len(GetConfig().Twitch.Admins)),
	}

	for _, u := range GetConfig().Twitch.Admins {
		c.admins[u] = true
	}

	d, err := ioutil.ReadFile(GetConfig().Twitch.ChannelListPath)
	if err != nil {
		log.Fatalf("unable to read channels %s", err)
	}
	if err := json.Unmarshal(d, &c.channels); err != nil {
		log.Fatalf("unable to read channels %s", err)
	}

	return c
}

// Connect open ws connection
func (c *TwitchChat) Connect() {
	var err error
	c.Lock()
	c.conn, _, err = c.dialer.Dial(GetConfig().Twitch.SocketURL, c.headers)
	c.Unlock()
	if err != nil {
		log.Printf("error connecting to twitch ws %s", err)
		c.reconnect()
	}

	c.send("PASS " + GetConfig().Twitch.OAuth)
	c.send("NICK " + GetConfig().Twitch.Nick)

	for _, ch := range c.channels {
		c.Join(ch, true)
	}
}

func (c *TwitchChat) reconnect() {
	if c.conn != nil {
		c.Lock()
		c.conn.Close()

		for ch, mc := range c.messages {
			close(mc)
			delete(c.messages, ch)
		}
		c.Unlock()
	}

	time.Sleep(SocketReconnectDelay)
	c.Connect()
}

// Run connect and start message read loop
func (c *TwitchChat) Run() {
	c.Connect()

	messagePattern := regexp.MustCompile(`:(.+)\!.+tmi\.twitch\.tv PRIVMSG #([a-z0-9_-]+) :(.+)`)

	for {
		err := c.conn.SetReadDeadline(time.Now().Add(SocketReadTimeout))
		if err != nil {
			c.reconnect()
			continue
		}

		c.RLock()
		_, msg, err := c.conn.ReadMessage()
		c.RUnlock()
		if err != nil {
			log.Printf("error reading message %s", err)
			c.reconnect()
			continue
		}

		if strings.Index(string(msg), "PING") == 0 {
			c.send(strings.Replace(string(msg), "PING", "PONG", -1))
			continue
		}

		l := messagePattern.FindAllStringSubmatch(string(msg), -1)
		for _, v := range l {
			c.RLock()
			mc, ok := c.messages[strings.ToLower(v[2])]
			c.RUnlock()
			if !ok {
				continue
			}

			data := strings.TrimSpace(v[3])
			data = strings.Replace(data, "ACTION", "/me", -1)
			data = strings.Replace(data, "", "", -1)
			m := &Message{
				Command: "MSG",
				Nick:    v[1],
				Data:    data,
				Time:    time.Now().UTC(),
			}

			c.runCommand(v[2], m)

			select {
			case mc <- m:
			default:
			}
		}
	}
}

func (c *TwitchChat) runCommand(source string, m *Message) {
	if _, ok := c.admins[m.Nick]; ok && m.Command == "MSG" {
		d := strings.Split(m.Data, " ")
		ld := strings.Split(strings.ToLower(m.Data), " ")

		if strings.EqualFold(d[0], "!join") {
			if err := c.Join(ld[1], false); err == nil {
				c.send("PRIVMSG #" + source + " :Logging " + ld[1])
			} else {
				c.send("PRIVMSG #" + source + " :Already logging " + ld[1])
			}
		} else if strings.EqualFold(d[0], "!leave") {
			if err := c.Leave(ld[1]); err == nil {
				c.send("PRIVMSG #" + source + " :Leaving " + ld[1])
			} else {
				c.send("PRIVMSG #" + source + " :Not logging " + ld[1])
			}
		} else if strings.EqualFold(d[0], "!channels") {
			c.send("PRIVMSG #" + source + " :Logging " + strings.Join(c.channels, ", "))
		}
	}
}

func (c *TwitchChat) send(m string) {
	c.writeLock.Lock()
	c.RLock()
	err := c.conn.WriteMessage(websocket.TextMessage, []byte(m+"\r\n"))
	c.RUnlock()
	if err == nil {
		time.Sleep(SocketWriteDebounce)
	}
	c.writeLock.Unlock()
	if err != nil {
		log.Printf("error sending message %s", err)
		c.reconnect()
	}
}

// Join channel
func (c *TwitchChat) Join(ch string, init bool) error {
	ch = strings.ToLower(ch)
	c.Lock()
	_, ok := c.messages[ch]
	if !ok {
		c.messages[ch] = make(chan *Message, MessageBufferSize)
	}
	c.Unlock()
	if ok {
		return ErrTwitchAlreadyInChannel
	}
	c.send("JOIN #" + ch)
	c.Lock()
	if messages, ok := c.messages[ch]; ok {
		go c.joinHandler(ch, messages)
	}
	c.Unlock()
	if init {
		return nil
	}
	return c.saveChannels()
}

// Leave channel
func (c *TwitchChat) Leave(ch string) error {
	ch = strings.ToLower(ch)
	c.Lock()
	_, ok := c.messages[ch]
	c.Unlock()
	if !ok {
		return ErrTwitchNotInChannel
	}
	c.send("PART #" + ch)
	c.Lock()
	delete(c.messages, ch)
	c.Unlock()
	return c.saveChannels()
}

func (c *TwitchChat) saveChannels() error {
	c.Lock()
	defer c.Unlock()
	c.channels = c.channels[:]
	for ch := range c.messages {
		c.channels = append(c.channels, ch)
	}
	f, err := os.Create(GetConfig().Twitch.ChannelListPath)
	if err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	data, err := json.Marshal(c.channels)
	if err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	var buf bytes.Buffer
	if err := json.Indent(&buf, data, "", "\t"); err != nil {
		log.Printf("error saving channel list %s", err)
	}
	f.Write(buf.Bytes())
	f.Close()
	return nil
}
<commit_msg>fixed duplicates in channels.json bug<commit_after>package common

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// errors
var (
	ErrTwitchAlreadyInChannel = errors.New("already in channel")
	ErrTwitchNotInChannel     = errors.New("not in channel")
)

// TwitchChat twitch chat client
type TwitchChat struct {
	sync.RWMutex
	conn        *websocket.Conn
	dialer      websocket.Dialer
	headers     http.Header
	messages    map[string]chan *Message
	channels    []string
	writeLock   sync.Mutex
	joinHandler TwitchJoinHandler
	admins      map[string]bool
}

// TwitchJoinHandler called when joining channels
type TwitchJoinHandler func(string, chan *Message)

// NewTwitchChat new twitch chat client
func NewTwitchChat(j TwitchJoinHandler) *TwitchChat {
	c := &TwitchChat{
		dialer:      websocket.Dialer{HandshakeTimeout: SocketHandshakeTimeout},
		headers:     http.Header{"Origin": []string{GetConfig().Twitch.OriginURL}},
		messages:    make(map[string]chan *Message, 0),
		channels:    make([]string, 0),
		joinHandler: j,
		admins:      make(map[string]bool, len(GetConfig().Twitch.Admins)),
	}

	for _, u := range GetConfig().Twitch.Admins {
		c.admins[u] = true
	}

	d, err := ioutil.ReadFile(GetConfig().Twitch.ChannelListPath)
	if err != nil {
		log.Fatalf("unable to read channels %s", err)
	}
	if err := json.Unmarshal(d, &c.channels); err != nil {
		log.Fatalf("unable to read channels %s", err)
	}

	return c
}

// Connect open ws connection
func (c *TwitchChat) Connect() {
	var err error
	c.Lock()
	c.conn, _, err = c.dialer.Dial(GetConfig().Twitch.SocketURL, c.headers)
	c.Unlock()
	if err != nil {
		log.Printf("error connecting to twitch ws %s", err)
		c.reconnect()
	}

	c.send("PASS " + GetConfig().Twitch.OAuth)
	c.send("NICK " + GetConfig().Twitch.Nick)

	for _, ch := range c.channels {
		c.Join(ch, true)
	}
}

func (c *TwitchChat) reconnect() {
	if c.conn != nil {
		c.Lock()
		c.conn.Close()

		for ch, mc := range c.messages {
			close(mc)
			delete(c.messages, ch)
		}
		c.Unlock()
	}

	time.Sleep(SocketReconnectDelay)
	c.Connect()
}

// Run connect and start message read loop
func (c *TwitchChat) Run() {
	c.Connect()

	messagePattern := regexp.MustCompile(`:(.+)\!.+tmi\.twitch\.tv PRIVMSG #([a-z0-9_-]+) :(.+)`)

	for {
		err := c.conn.SetReadDeadline(time.Now().Add(SocketReadTimeout))
		if err != nil {
			c.reconnect()
			continue
		}

		c.RLock()
		_, msg, err := c.conn.ReadMessage()
		c.RUnlock()
		if err != nil {
			log.Printf("error reading message %s", err)
			c.reconnect()
			continue
		}

		if strings.Index(string(msg), "PING") == 0 {
			c.send(strings.Replace(string(msg), "PING", "PONG", -1))
			continue
		}

		l := messagePattern.FindAllStringSubmatch(string(msg), -1)
		for _, v := range l {
			c.RLock()
			mc, ok := c.messages[strings.ToLower(v[2])]
			c.RUnlock()
			if !ok {
				continue
			}

			data := strings.TrimSpace(v[3])
			data = strings.Replace(data, "ACTION", "/me", -1)
			data = strings.Replace(data, "", "", -1)
			m := &Message{
				Command: "MSG",
				Nick:    v[1],
				Data:    data,
				Time:    time.Now().UTC(),
			}

			c.runCommand(v[2], m)

			select {
			case mc <- m:
			default:
			}
		}
	}
}

func (c *TwitchChat) runCommand(source string, m *Message) {
	if _, ok := c.admins[m.Nick]; ok && m.Command == "MSG" {
		d := strings.Split(m.Data, " ")
		ld := strings.Split(strings.ToLower(m.Data), " ")

		if strings.EqualFold(d[0], "!join") {
			if err := c.Join(ld[1], false); err == nil {
				c.send("PRIVMSG #" + source + " :Logging " + ld[1])
			} else {
				c.send("PRIVMSG #" + source + " :Already logging " + ld[1])
			}
		} else if strings.EqualFold(d[0], "!leave") {
			if err := c.Leave(ld[1]); err == nil {
				c.send("PRIVMSG #" + source + " :Leaving " + ld[1])
			} else {
				c.send("PRIVMSG #" + source + " :Not logging " + ld[1])
			}
		} else if strings.EqualFold(d[0], "!channels") {
			c.send("PRIVMSG #" + source + " :Logging " + strings.Join(c.channels, ", "))
		}
	}
}

func (c *TwitchChat) send(m string) {
	c.writeLock.Lock()
	c.RLock()
	err := c.conn.WriteMessage(websocket.TextMessage, []byte(m+"\r\n"))
	c.RUnlock()
	if err == nil {
		time.Sleep(SocketWriteDebounce)
	}
	c.writeLock.Unlock()
	if err != nil {
		log.Printf("error sending message %s", err)
		c.reconnect()
	}
}

// Join channel
func (c *TwitchChat) Join(ch string, init bool) error {
	ch = strings.ToLower(ch)
	c.Lock()
	_, ok := c.messages[ch]
	if !ok {
		c.messages[ch] = make(chan *Message, MessageBufferSize)
	}
	c.Unlock()
	if ok {
		return ErrTwitchAlreadyInChannel
	}
	c.send("JOIN #" + ch)
	c.Lock()
	if messages, ok := c.messages[ch]; ok {
		go c.joinHandler(ch, messages)
	}
	c.Unlock()
	if init {
		return nil
	}
	return c.saveChannels()
}

// Leave channel
func (c *TwitchChat) Leave(ch string) error {
	ch = strings.ToLower(ch)
	c.Lock()
	_, ok := c.messages[ch]
	c.Unlock()
	if !ok {
		return ErrTwitchNotInChannel
	}
	c.send("PART #" + ch)
	c.Lock()
	delete(c.messages, ch)
	c.Unlock()
	return c.saveChannels()
}

func (c *TwitchChat) saveChannels() error {
	c.Lock()
	defer c.Unlock()
	c.channels = []string{}
	for ch := range c.messages {
		c.channels = append(c.channels, ch)
	}
	f, err := os.Create(GetConfig().Twitch.ChannelListPath)
	if err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	data, err := json.Marshal(c.channels)
	if err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	var buf bytes.Buffer
	if err := json.Indent(&buf, data, "", "\t"); err != nil {
		log.Printf("error saving channel list %s", err)
	}
	f.Write(buf.Bytes())
	f.Close()
	return nil
}
<|endoftext|>


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

package declarative

import (
	"github.com/lxn/walk"
)

type TreeView struct {
	AssignTo             **walk.TreeView
	Name                 string
	Enabled              Property
	Visible              Property
	Font                 Font
	ToolTipText          Property
	MinSize              Size
	MaxSize              Size
	StretchFactor        int
	Row                  int
	RowSpan              int
	Column               int
	ColumnSpan           int
	ContextMenuActions   []*walk.Action
	OnKeyDown            walk.KeyEventHandler
	OnMouseDown          walk.MouseEventHandler
	OnMouseMove          walk.MouseEventHandler
	OnMouseUp            walk.MouseEventHandler
	OnSizeChanged        walk.EventHandler
	Model                walk.TreeModel
	OnCurrentItemChanged walk.EventHandler
	OnItemCollapsing     walk.TreeItemEventHandler
	OnItemCollapsed      walk.TreeItemEventHandler
	OnItemExpanding      walk.TreeItemEventHandler
	OnItemExpanded       walk.TreeItemEventHandler
}

func (tv TreeView) Create(builder *Builder) error {
	w, err := walk.NewTreeView(builder.Parent())
	if err != nil {
		return err
	}

	return builder.InitWidget(tv, w, func() error {
		if err := w.SetModel(tv.Model); err != nil {
			return err
		}

		if tv.OnCurrentItemChanged != nil {
			w.CurrentItemChanged().Attach(tv.OnCurrentItemChanged)
		}

		if tv.OnItemCollapsing != nil {
			w.ItemCollapsing().Attach(tv.OnItemCollapsing)
		}

		if tv.OnItemCollapsed != nil {
			w.ItemCollapsed().Attach(tv.OnItemCollapsed)
		}

		if tv.OnItemExpanding != nil {
			w.ItemExpanding().Attach(tv.OnItemExpanding)
		}

		if tv.OnItemExpanded != nil {
			w.ItemExpanded().Attach(tv.OnItemExpanded)
		}

		if tv.AssignTo != nil {
			*tv.AssignTo = w
		}

		return nil
	})
}

func (w TreeView) WidgetInfo() (name string, disabled, hidden bool, font *Font, toolTipText string, minSize, maxSize Size, stretchFactor, row, rowSpan, column, columnSpan int, contextMenuActions []*walk.Action, OnKeyDown walk.KeyEventHandler, OnMouseDown walk.MouseEventHandler, OnMouseMove walk.MouseEventHandler, OnMouseUp walk.MouseEventHandler, OnSizeChanged walk.EventHandler) {
	return w.Name, false, false, &w.Font, "", w.MinSize, w.MaxSize, w.StretchFactor, w.Row, w.RowSpan, w.Column, w.ColumnSpan, w.ContextMenuActions, w.OnKeyDown, w.OnMouseDown, w.OnMouseMove, w.OnMouseUp, w.OnSizeChanged
}
<commit_msg>Revert "declarative/TreeView: Add OnItemExpanding and OnItemCollapsing"<commit_after>// Copyright 2012 The Walk Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package declarative

import (
	"github.com/lxn/walk"
)

type TreeView struct {
	AssignTo             **walk.TreeView
	Name                 string
	Enabled              Property
	Visible              Property
	Font                 Font
	ToolTipText          Property
	MinSize              Size
	MaxSize              Size
	StretchFactor        int
	Row                  int
	RowSpan              int
	Column               int
	ColumnSpan           int
	ContextMenuActions   []*walk.Action
	OnKeyDown            walk.KeyEventHandler
	OnMouseDown          walk.MouseEventHandler
	OnMouseMove          walk.MouseEventHandler
	OnMouseUp            walk.MouseEventHandler
	OnSizeChanged        walk.EventHandler
	Model                walk.TreeModel
	OnCurrentItemChanged walk.EventHandler
	OnItemCollapsed      walk.TreeItemEventHandler
	OnItemExpanded       walk.TreeItemEventHandler
}

func (tv TreeView) Create(builder *Builder) error {
	w, err := walk.NewTreeView(builder.Parent())
	if err != nil {
		return err
	}

	return builder.InitWidget(tv, w, func() error {
		if err := w.SetModel(tv.Model); err != nil {
			return err
		}

		if tv.OnCurrentItemChanged != nil {
			w.CurrentItemChanged().Attach(tv.OnCurrentItemChanged)
		}

		if tv.OnItemCollapsed != nil {
			w.ItemCollapsed().Attach(tv.OnItemCollapsed)
		}

		if tv.OnItemExpanded != nil {
			w.ItemExpanded().Attach(tv.OnItemExpanded)
		}

		if tv.AssignTo != nil {
			*tv.AssignTo = w
		}

		return nil
	})
}

func (w TreeView) WidgetInfo() (name string, disabled, hidden bool, font *Font, toolTipText string, minSize, maxSize Size, stretchFactor, row, rowSpan, column, columnSpan int, contextMenuActions []*walk.Action, OnKeyDown walk.KeyEventHandler, OnMouseDown walk.MouseEventHandler, OnMouseMove walk.MouseEventHandler, OnMouseUp walk.MouseEventHandler, OnSizeChanged walk.EventHandler) {
	return w.Name, false, false, &w.Font, "", w.MinSize, w.MaxSize, w.StretchFactor, w.Row, w.RowSpan, w.Column, w.ColumnSpan, w.ContextMenuActions, w.OnKeyDown, w.OnMouseDown, w.OnMouseMove, w.OnMouseUp, w.OnSizeChanged
}
<|endoftext|>


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

package client

import (
	"fmt"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/keybase/cli"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/service"
)

// GetExtraFlags gets the extra fork-related flags for this platform
func GetExtraFlags() []cli.Flag {
	return []cli.Flag{
		cli.BoolFlag{
			Name:  "no-auto-fork, F",
			Usage: "disable auto-fork of background service",
		},
		cli.BoolFlag{
			Name:  "auto-fork",
			Usage: "enable auto-fork of background service",
		},
	}
}

// ForkServerNix forks a new background Keybase service, and waits until it's
// pingable. It will only do something useful on Unixes; it won't work on
// Windows (probably?). Returns an error if anything bad happens; otherwise,
// assume that the server was successfully started up.
func ForkServerNix(cl libkb.CommandLine) error {
	srv := service.NewService(true /* isDaemon */)

	// If we try to get an exclusive lock and succeed, it means we
	// need to relaunch the daemon since it's dead
	G.Log.Debug("Getting flock")
	err := srv.GetExclusiveLock()
	if err == nil {
		G.Log.Debug("Flocked! Server must have died")
		srv.ReleaseLock()
		err = spawnServer(cl)
		if err != nil {
			G.Log.Errorf("Error in spawning server process: %s", err)
			return err
		}
		err = pingLoop()
		if err != nil {
			G.Log.Errorf("Ping failure after server fork: %s", err)
			return err
		}
	} else {
		G.Log.Debug("The server is still up")
		err = nil
	}

	return err
}

func pingLoop() error {
	var err error
	for i := 0; i < 10; i++ {
		_, _, err = G.GetSocket()
		if err == nil {
			G.Log.Debug("Connected (%d)", i)
			return nil
		}
		G.ClearSocketError()
		G.Log.Debug("Failed to connect to socket (%d): %s", i, err)
		err = nil
		time.Sleep(200 * time.Millisecond)
	}
	return nil
}

func makeServerCommandLine(cl libkb.CommandLine) (arg0 string, args []string, err error) {
	// ForkExec requires an absolute path to the binary. LookPath() gets this
	// for us, or correctly leaves arg0 alone if it's already a path.
	arg0, err = exec.LookPath(os.Args[0])
	if err != nil {
		return
	}

	// Fixme: This isn't ideal, it would be better to specify when the args
	// are defined if they should be reexported to the server, and if so, then
	// we should automate the reconstruction of the argument vector.  Let's do
	// this when we yank out keybase/cli
	bools := []string{
		"debug",
		"api-dump-unsafe",
		"plain-logging",
	}

	strings := []string{
		"home",
		"server",
		"config",
		"session",
		"proxy",
		"username",
		"gpg-home",
		"gpg",
		"secret-keyring",
		"pid-file",
		"socket-file",
		"gpg-options",
		"local-rpc-debug-unsafe",
	}
	args = append(args, arg0)

	for _, b := range bools {
		if isSet, isTrue := cl.GetBool(b, true); isSet && isTrue {
			args = append(args, "--"+b)
		}
	}

	for _, s := range strings {
		if v := cl.GetGString(s); len(v) > 0 {
			args = append(args, "--"+s, v)
		}
	}

	args = append(args, "service")

	var chdir string
	chdir, err = G.Env.GetServiceSpawnDir()
	if err != nil {
		return
	}

	G.Log.Info("| Setting run directory for keybase service to %s", chdir)
	args = append(args, "--chdir", chdir)

	G.Log.Debug("| Made server args: %s %v", arg0, args)

	return
}

func spawnServer(cl libkb.CommandLine) (err error) {

	var files []uintptr
	var cmd string
	var args []string
	var devnull, log *os.File
	var pid int

	defer func() {
		if err != nil {
			if devnull != nil {
				devnull.Close()
			}
			if log != nil {
				log.Close()
			}
		}
	}()

	if devnull, err = os.OpenFile("/dev/null", os.O_RDONLY, 0); err != nil {
		return
	}
	files = append(files, devnull.Fd())

	if G.Env.GetSplitLogOutput() {
		files = append(files, uintptr(1), uintptr(2))
	} else {
		if _, log, err = libkb.OpenLogFile(); err != nil {
			return
		}
		files = append(files, log.Fd(), log.Fd())
	}

	attr := syscall.ProcAttr{
		Env:   os.Environ(),
		Sys:   &syscall.SysProcAttr{Setsid: true},
		Files: files,
	}

	cmd, args, err = makeServerCommandLine(cl)
	if err != nil {
		return err
	}

	pid, err = syscall.ForkExec(cmd, args, &attr)
	if err != nil {
		err = fmt.Errorf("Error in ForkExec: %s", err)
	} else {
		G.Log.Info("Forking background server with pid=%d", pid)
	}
	return err
}
<commit_msg>Usage, description should be sentences.<commit_after>// +build darwin dragonfly freebsd linux netbsd openbsd solaris

package client

import (
	"fmt"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/keybase/cli"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/service"
)

// GetExtraFlags gets the extra fork-related flags for this platform
func GetExtraFlags() []cli.Flag {
	return []cli.Flag{
		cli.BoolFlag{
			Name:  "no-auto-fork, F",
			Usage: "Disable auto-fork of background service.",
		},
		cli.BoolFlag{
			Name:  "auto-fork",
			Usage: "Enable auto-fork of background service.",
		},
	}
}

// ForkServerNix forks a new background Keybase service, and waits until it's
// pingable. It will only do something useful on Unixes; it won't work on
// Windows (probably?). Returns an error if anything bad happens; otherwise,
// assume that the server was successfully started up.
func ForkServerNix(cl libkb.CommandLine) error {
	srv := service.NewService(true /* isDaemon */)

	// If we try to get an exclusive lock and succeed, it means we
	// need to relaunch the daemon since it's dead
	G.Log.Debug("Getting flock")
	err := srv.GetExclusiveLock()
	if err == nil {
		G.Log.Debug("Flocked! Server must have died")
		srv.ReleaseLock()
		err = spawnServer(cl)
		if err != nil {
			G.Log.Errorf("Error in spawning server process: %s", err)
			return err
		}
		err = pingLoop()
		if err != nil {
			G.Log.Errorf("Ping failure after server fork: %s", err)
			return err
		}
	} else {
		G.Log.Debug("The server is still up")
		err = nil
	}

	return err
}

func pingLoop() error {
	var err error
	for i := 0; i < 10; i++ {
		_, _, err = G.GetSocket()
		if err == nil {
			G.Log.Debug("Connected (%d)", i)
			return nil
		}
		G.ClearSocketError()
		G.Log.Debug("Failed to connect to socket (%d): %s", i, err)
		err = nil
		time.Sleep(200 * time.Millisecond)
	}
	return nil
}

func makeServerCommandLine(cl libkb.CommandLine) (arg0 string, args []string, err error) {
	// ForkExec requires an absolute path to the binary. LookPath() gets this
	// for us, or correctly leaves arg0 alone if it's already a path.
	arg0, err = exec.LookPath(os.Args[0])
	if err != nil {
		return
	}

	// Fixme: This isn't ideal, it would be better to specify when the args
	// are defined if they should be reexported to the server, and if so, then
	// we should automate the reconstruction of the argument vector.  Let's do
	// this when we yank out keybase/cli
	bools := []string{
		"debug",
		"api-dump-unsafe",
		"plain-logging",
	}

	strings := []string{
		"home",
		"server",
		"config",
		"session",
		"proxy",
		"username",
		"gpg-home",
		"gpg",
		"secret-keyring",
		"pid-file",
		"socket-file",
		"gpg-options",
		"local-rpc-debug-unsafe",
	}
	args = append(args, arg0)

	for _, b := range bools {
		if isSet, isTrue := cl.GetBool(b, true); isSet && isTrue {
			args = append(args, "--"+b)
		}
	}

	for _, s := range strings {
		if v := cl.GetGString(s); len(v) > 0 {
			args = append(args, "--"+s, v)
		}
	}

	args = append(args, "service")

	var chdir string
	chdir, err = G.Env.GetServiceSpawnDir()
	if err != nil {
		return
	}

	G.Log.Info("| Setting run directory for keybase service to %s", chdir)
	args = append(args, "--chdir", chdir)

	G.Log.Debug("| Made server args: %s %v", arg0, args)

	return
}

func spawnServer(cl libkb.CommandLine) (err error) {

	var files []uintptr
	var cmd string
	var args []string
	var devnull, log *os.File
	var pid int

	defer func() {
		if err != nil {
			if devnull != nil {
				devnull.Close()
			}
			if log != nil {
				log.Close()
			}
		}
	}()

	if devnull, err = os.OpenFile("/dev/null", os.O_RDONLY, 0); err != nil {
		return
	}
	files = append(files, devnull.Fd())

	if G.Env.GetSplitLogOutput() {
		files = append(files, uintptr(1), uintptr(2))
	} else {
		if _, log, err = libkb.OpenLogFile(); err != nil {
			return
		}
		files = append(files, log.Fd(), log.Fd())
	}

	attr := syscall.ProcAttr{
		Env:   os.Environ(),
		Sys:   &syscall.SysProcAttr{Setsid: true},
		Files: files,
	}

	cmd, args, err = makeServerCommandLine(cl)
	if err != nil {
		return err
	}

	pid, err = syscall.ForkExec(cmd, args, &attr)
	if err != nil {
		err = fmt.Errorf("Error in ForkExec: %s", err)
	} else {
		G.Log.Info("Forking background server with pid=%d", pid)
	}
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"strings"
	"time"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	mproto "github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/vt/dbconfigs"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	"github.com/youtube/vitess/go/vt/tabletmanager"
	"github.com/youtube/vitess/go/vt/tabletserver"
	tproto "github.com/youtube/vitess/go/vt/tabletserver/proto"
	"github.com/youtube/vitess/go/vt/tabletserver/tabletconn"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/topo/topoproto"
	"github.com/youtube/vitess/go/vt/wrangler"

	pbq "github.com/youtube/vitess/go/vt/proto/query"
	pb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// tablet contains all the data for an individual tablet.
type tablet struct {
	// configuration parameters
	keyspace   string
	shard      string
	tabletType pb.TabletType
	dbname     string

	// objects built at construction time
	qsc   tabletserver.Controller
	agent *tabletmanager.ActionAgent
}

// tabletMap maps the tablet uid to the tablet record
var tabletMap map[uint32]*tablet

// initTabletMap creates the action agents and associated data structures
// for all tablets
func initTabletMap(ts topo.Server, topology string, mysqld mysqlctl.MysqlDaemon, dbcfgs *dbconfigs.DBConfigs, mycnf *mysqlctl.Mycnf) {
	tabletMap = make(map[uint32]*tablet)

	ctx := context.Background()

	// disable publishing of stats from query service
	flag.Lookup("queryserver-config-enable-publish-stats").Value.Set("false")

	var uid uint32 = 1
	keyspaceMap := make(map[string]bool)
	for _, entry := range strings.Split(topology, ",") {
		slash := strings.IndexByte(entry, '/')
		column := strings.IndexByte(entry, ':')
		if slash == -1 || column == -1 {
			log.Fatalf("invalid topology entry: %v", entry)
		}

		keyspace := entry[:slash]
		shard := entry[slash+1 : column]
		dbname := entry[column+1:]
		keyspaceMap[keyspace] = true

		localDBConfigs := &(*dbcfgs)
		localDBConfigs.App.DbName = dbname

		// create the master
		alias := &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating master tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		masterController := tabletserver.NewServer()
		masterAgent := tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), masterController, localDBConfigs, mysqld, keyspace, shard, dbname, "replica")
		if err := masterAgent.TabletExternallyReparented(ctx, ""); err != nil {
			log.Fatalf("TabletExternallyReparented failed on master: %v", err)
		}
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_MASTER,
			dbname:     dbname,

			qsc:   masterController,
			agent: masterAgent,
		}
		uid++

		// create a replica slave
		alias = &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating replica tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		replicaController := tabletserver.NewServer()
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_REPLICA,
			dbname:     dbname,

			qsc:   replicaController,
			agent: tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), replicaController, localDBConfigs, mysqld, keyspace, shard, dbname, "replica"),
		}
		uid++

		// create a rdonly slave
		alias = &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating rdonly tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		rdonlyController := tabletserver.NewServer()
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_RDONLY,
			dbname:     dbname,

			qsc:   rdonlyController,
			agent: tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), rdonlyController, localDBConfigs, mysqld, keyspace, shard, dbname, "rdonly"),
		}
		uid++
	}

	// Rebuild the SrvKeyspace objects, we we can support range-based
	// sharding queries.
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, nil, 30*time.Second /*lockTimeout*/)
	for keyspace := range keyspaceMap {
		if err := wr.RebuildKeyspaceGraph(ctx, keyspace, nil, true); err != nil {
			log.Fatalf("cannot rebuild %v: %v", keyspace, err)
		}
	}

	// Register the tablet dialer
	tabletconn.RegisterDialer("internal", dialer)
	*tabletconn.TabletProtocol = "internal"
}

// dialer is our tabletconn.Dialer
func dialer(ctx context.Context, endPoint *pb.EndPoint, keyspace, shard string, tabletType pb.TabletType, timeout time.Duration) (tabletconn.TabletConn, error) {
	tablet, ok := tabletMap[endPoint.Uid]
	if !ok {
		return nil, tabletconn.OperationalError("connection refused")
	}

	return &internalTabletConn{
		tablet:   tablet,
		endPoint: endPoint,
	}, nil
}

// internalTabletConn implements tabletconn.TabletConn by forwarding everything
// to the tablet
type internalTabletConn struct {
	tablet   *tablet
	endPoint *pb.EndPoint
}

// Execute is part of tabletconn.TabletConn
func (itc *internalTabletConn) Execute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (*mproto.QueryResult, error) {
	reply := &mproto.QueryResult{}
	if err := itc.tablet.qsc.QueryService().Execute(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Query{
		Sql:           query,
		BindVariables: bindVars,
		TransactionId: transactionID,
	}, reply); err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return reply, nil
}

// ExecuteBatch is part of tabletconn.TabletConn
func (itc *internalTabletConn) ExecuteBatch(ctx context.Context, queries []tproto.BoundQuery, asTransaction bool, transactionID int64) (*tproto.QueryResultList, error) {
	reply := &tproto.QueryResultList{}
	if err := itc.tablet.qsc.QueryService().ExecuteBatch(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.QueryList{
		Queries:       queries,
		AsTransaction: asTransaction,
		TransactionId: transactionID,
	}, reply); err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return reply, nil
}

// StreamExecute is part of tabletconn.TabletConn
func (itc *internalTabletConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *mproto.QueryResult, tabletconn.ErrFunc, error) {
	result := make(chan *mproto.QueryResult, 10)
	var finalErr error

	go func() {
		finalErr = itc.tablet.qsc.QueryService().StreamExecute(ctx, &pbq.Target{
			Keyspace:   itc.tablet.keyspace,
			Shard:      itc.tablet.shard,
			TabletType: itc.tablet.tabletType,
		}, &tproto.Query{
			Sql:           query,
			BindVariables: bindVars,
			TransactionId: transactionID,
		}, func(reply *mproto.QueryResult) error {
			result <- reply
			return nil
		})

		// the client will only access finalErr after the
		// channel is closed, and then it's already set.
		close(result)
	}()

	return result, func() error {
		return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(finalErr))
	}, nil
}

// Begin is part of tabletconn.TabletConn
func (itc *internalTabletConn) Begin(ctx context.Context) (transactionID int64, err error) {
	result := &tproto.TransactionInfo{}
	if err := itc.tablet.qsc.QueryService().Begin(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Session{}, result); err != nil {
		return 0, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return result.TransactionId, nil
}

// Commit is part of tabletconn.TabletConn
func (itc *internalTabletConn) Commit(ctx context.Context, transactionID int64) error {
	if err := itc.tablet.qsc.QueryService().Commit(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Session{
		TransactionId: transactionID,
	}); err != nil {
		return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return nil
}

// Rollback is part of tabletconn.TabletConn
func (itc *internalTabletConn) Rollback(ctx context.Context, transactionID int64) error {
	if err := itc.tablet.qsc.QueryService().Rollback(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Session{
		TransactionId: transactionID,
	}); err != nil {
		return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return nil
}

// Execute2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) Execute2(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (*mproto.QueryResult, error) {
	return itc.Execute(ctx, query, bindVars, transactionID)
}

// ExecuteBatch2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) ExecuteBatch2(ctx context.Context, queries []tproto.BoundQuery, asTransaction bool, transactionID int64) (*tproto.QueryResultList, error) {
	return itc.ExecuteBatch(ctx, queries, asTransaction, transactionID)
}

// Begin2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) Begin2(ctx context.Context) (transactionID int64, err error) {
	return itc.Begin(ctx)
}

// Commit2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) Commit2(ctx context.Context, transactionID int64) error {
	return itc.Commit(ctx, transactionID)
}

// Rollback2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) Rollback2(ctx context.Context, transactionID int64) error {
	return itc.Rollback(ctx, transactionID)
}

// StreamExecute2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) StreamExecute2(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *mproto.QueryResult, tabletconn.ErrFunc, error) {
	return itc.StreamExecute(ctx, query, bindVars, transactionID)
}

// Close is part of tabletconn.TabletConn
func (itc *internalTabletConn) Close() {
}

// SetTarget is part of tabletconn.TabletConn
func (itc *internalTabletConn) SetTarget(keyspace, shard string, tabletType pb.TabletType) error {
	return nil
}

// EndPoint is part of tabletconn.TabletConn
func (itc *internalTabletConn) EndPoint() *pb.EndPoint {
	return itc.endPoint
}

// SplitQuery is part of tabletconn.TabletConn
func (itc *internalTabletConn) SplitQuery(ctx context.Context, query tproto.BoundQuery, splitColumn string, splitCount int) ([]tproto.QuerySplit, error) {
	reply := &tproto.SplitQueryResult{}
	if err := itc.tablet.qsc.QueryService().SplitQuery(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.SplitQueryRequest{
		Query:       query,
		SplitColumn: splitColumn,
		SplitCount:  splitCount,
	}, reply); err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return reply.Queries, nil
}

// StreamHealth is part of tabletconn.TabletConn
func (itc *internalTabletConn) StreamHealth(ctx context.Context) (<-chan *pbq.StreamHealthResponse, tabletconn.ErrFunc, error) {
	result := make(chan *pbq.StreamHealthResponse, 10)

	id, err := itc.tablet.qsc.QueryService().StreamHealthRegister(result)
	if err != nil {
		return nil, nil, err
	}

	var finalErr error
	go func() {
		select {
		case <-ctx.Done():
		}

		finalErr = itc.tablet.qsc.QueryService().StreamHealthUnregister(id)
		close(result)
	}()

	return result, func() error {
		return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(finalErr))
	}, nil
}
<commit_msg>Fixing a case where err could be nil and trigger a panic.<commit_after>package main

import (
	"flag"
	"fmt"
	"strings"
	"time"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	mproto "github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/vt/dbconfigs"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	"github.com/youtube/vitess/go/vt/tabletmanager"
	"github.com/youtube/vitess/go/vt/tabletserver"
	tproto "github.com/youtube/vitess/go/vt/tabletserver/proto"
	"github.com/youtube/vitess/go/vt/tabletserver/tabletconn"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/topo/topoproto"
	"github.com/youtube/vitess/go/vt/wrangler"

	pbq "github.com/youtube/vitess/go/vt/proto/query"
	pb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// tablet contains all the data for an individual tablet.
type tablet struct {
	// configuration parameters
	keyspace   string
	shard      string
	tabletType pb.TabletType
	dbname     string

	// objects built at construction time
	qsc   tabletserver.Controller
	agent *tabletmanager.ActionAgent
}

// tabletMap maps the tablet uid to the tablet record
var tabletMap map[uint32]*tablet

// initTabletMap creates the action agents and associated data structures
// for all tablets
func initTabletMap(ts topo.Server, topology string, mysqld mysqlctl.MysqlDaemon, dbcfgs *dbconfigs.DBConfigs, mycnf *mysqlctl.Mycnf) {
	tabletMap = make(map[uint32]*tablet)

	ctx := context.Background()

	// disable publishing of stats from query service
	flag.Lookup("queryserver-config-enable-publish-stats").Value.Set("false")

	var uid uint32 = 1
	keyspaceMap := make(map[string]bool)
	for _, entry := range strings.Split(topology, ",") {
		slash := strings.IndexByte(entry, '/')
		column := strings.IndexByte(entry, ':')
		if slash == -1 || column == -1 {
			log.Fatalf("invalid topology entry: %v", entry)
		}

		keyspace := entry[:slash]
		shard := entry[slash+1 : column]
		dbname := entry[column+1:]
		keyspaceMap[keyspace] = true

		localDBConfigs := &(*dbcfgs)
		localDBConfigs.App.DbName = dbname

		// create the master
		alias := &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating master tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		masterController := tabletserver.NewServer()
		masterAgent := tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), masterController, localDBConfigs, mysqld, keyspace, shard, dbname, "replica")
		if err := masterAgent.TabletExternallyReparented(ctx, ""); err != nil {
			log.Fatalf("TabletExternallyReparented failed on master: %v", err)
		}
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_MASTER,
			dbname:     dbname,

			qsc:   masterController,
			agent: masterAgent,
		}
		uid++

		// create a replica slave
		alias = &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating replica tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		replicaController := tabletserver.NewServer()
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_REPLICA,
			dbname:     dbname,

			qsc:   replicaController,
			agent: tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), replicaController, localDBConfigs, mysqld, keyspace, shard, dbname, "replica"),
		}
		uid++

		// create a rdonly slave
		alias = &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating rdonly tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		rdonlyController := tabletserver.NewServer()
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_RDONLY,
			dbname:     dbname,

			qsc:   rdonlyController,
			agent: tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), rdonlyController, localDBConfigs, mysqld, keyspace, shard, dbname, "rdonly"),
		}
		uid++
	}

	// Rebuild the SrvKeyspace objects, we we can support range-based
	// sharding queries.
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, nil, 30*time.Second /*lockTimeout*/)
	for keyspace := range keyspaceMap {
		if err := wr.RebuildKeyspaceGraph(ctx, keyspace, nil, true); err != nil {
			log.Fatalf("cannot rebuild %v: %v", keyspace, err)
		}
	}

	// Register the tablet dialer
	tabletconn.RegisterDialer("internal", dialer)
	*tabletconn.TabletProtocol = "internal"
}

// dialer is our tabletconn.Dialer
func dialer(ctx context.Context, endPoint *pb.EndPoint, keyspace, shard string, tabletType pb.TabletType, timeout time.Duration) (tabletconn.TabletConn, error) {
	tablet, ok := tabletMap[endPoint.Uid]
	if !ok {
		return nil, tabletconn.OperationalError("connection refused")
	}

	return &internalTabletConn{
		tablet:   tablet,
		endPoint: endPoint,
	}, nil
}

// internalTabletConn implements tabletconn.TabletConn by forwarding everything
// to the tablet
type internalTabletConn struct {
	tablet   *tablet
	endPoint *pb.EndPoint
}

// Execute is part of tabletconn.TabletConn
func (itc *internalTabletConn) Execute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (*mproto.QueryResult, error) {
	reply := &mproto.QueryResult{}
	if err := itc.tablet.qsc.QueryService().Execute(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Query{
		Sql:           query,
		BindVariables: bindVars,
		TransactionId: transactionID,
	}, reply); err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return reply, nil
}

// ExecuteBatch is part of tabletconn.TabletConn
func (itc *internalTabletConn) ExecuteBatch(ctx context.Context, queries []tproto.BoundQuery, asTransaction bool, transactionID int64) (*tproto.QueryResultList, error) {
	reply := &tproto.QueryResultList{}
	if err := itc.tablet.qsc.QueryService().ExecuteBatch(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.QueryList{
		Queries:       queries,
		AsTransaction: asTransaction,
		TransactionId: transactionID,
	}, reply); err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return reply, nil
}

// StreamExecute is part of tabletconn.TabletConn
func (itc *internalTabletConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *mproto.QueryResult, tabletconn.ErrFunc, error) {
	result := make(chan *mproto.QueryResult, 10)
	var finalErr error

	go func() {
		finalErr = itc.tablet.qsc.QueryService().StreamExecute(ctx, &pbq.Target{
			Keyspace:   itc.tablet.keyspace,
			Shard:      itc.tablet.shard,
			TabletType: itc.tablet.tabletType,
		}, &tproto.Query{
			Sql:           query,
			BindVariables: bindVars,
			TransactionId: transactionID,
		}, func(reply *mproto.QueryResult) error {
			result <- reply
			return nil
		})

		// the client will only access finalErr after the
		// channel is closed, and then it's already set.
		close(result)
	}()

	return result, func() error {
		if finalErr != nil {
			return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(finalErr))
		}
		return nil
	}, nil
}

// Begin is part of tabletconn.TabletConn
func (itc *internalTabletConn) Begin(ctx context.Context) (transactionID int64, err error) {
	result := &tproto.TransactionInfo{}
	if err := itc.tablet.qsc.QueryService().Begin(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Session{}, result); err != nil {
		return 0, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return result.TransactionId, nil
}

// Commit is part of tabletconn.TabletConn
func (itc *internalTabletConn) Commit(ctx context.Context, transactionID int64) error {
	if err := itc.tablet.qsc.QueryService().Commit(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Session{
		TransactionId: transactionID,
	}); err != nil {
		return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return nil
}

// Rollback is part of tabletconn.TabletConn
func (itc *internalTabletConn) Rollback(ctx context.Context, transactionID int64) error {
	if err := itc.tablet.qsc.QueryService().Rollback(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Session{
		TransactionId: transactionID,
	}); err != nil {
		return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return nil
}

// Execute2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) Execute2(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (*mproto.QueryResult, error) {
	return itc.Execute(ctx, query, bindVars, transactionID)
}

// ExecuteBatch2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) ExecuteBatch2(ctx context.Context, queries []tproto.BoundQuery, asTransaction bool, transactionID int64) (*tproto.QueryResultList, error) {
	return itc.ExecuteBatch(ctx, queries, asTransaction, transactionID)
}

// Begin2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) Begin2(ctx context.Context) (transactionID int64, err error) {
	return itc.Begin(ctx)
}

// Commit2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) Commit2(ctx context.Context, transactionID int64) error {
	return itc.Commit(ctx, transactionID)
}

// Rollback2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) Rollback2(ctx context.Context, transactionID int64) error {
	return itc.Rollback(ctx, transactionID)
}

// StreamExecute2 is part of tabletconn.TabletConn
func (itc *internalTabletConn) StreamExecute2(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *mproto.QueryResult, tabletconn.ErrFunc, error) {
	return itc.StreamExecute(ctx, query, bindVars, transactionID)
}

// Close is part of tabletconn.TabletConn
func (itc *internalTabletConn) Close() {
}

// SetTarget is part of tabletconn.TabletConn
func (itc *internalTabletConn) SetTarget(keyspace, shard string, tabletType pb.TabletType) error {
	return nil
}

// EndPoint is part of tabletconn.TabletConn
func (itc *internalTabletConn) EndPoint() *pb.EndPoint {
	return itc.endPoint
}

// SplitQuery is part of tabletconn.TabletConn
func (itc *internalTabletConn) SplitQuery(ctx context.Context, query tproto.BoundQuery, splitColumn string, splitCount int) ([]tproto.QuerySplit, error) {
	reply := &tproto.SplitQueryResult{}
	if err := itc.tablet.qsc.QueryService().SplitQuery(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.SplitQueryRequest{
		Query:       query,
		SplitColumn: splitColumn,
		SplitCount:  splitCount,
	}, reply); err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return reply.Queries, nil
}

// StreamHealth is part of tabletconn.TabletConn
func (itc *internalTabletConn) StreamHealth(ctx context.Context) (<-chan *pbq.StreamHealthResponse, tabletconn.ErrFunc, error) {
	result := make(chan *pbq.StreamHealthResponse, 10)

	id, err := itc.tablet.qsc.QueryService().StreamHealthRegister(result)
	if err != nil {
		return nil, nil, err
	}

	var finalErr error
	go func() {
		select {
		case <-ctx.Done():
		}

		finalErr = itc.tablet.qsc.QueryService().StreamHealthUnregister(id)
		close(result)
	}()

	return result, func() error {
		if finalErr != nil {
			return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(finalErr))
		}
		return nil
	}, nil
}
<|endoftext|>


<commit_before>package dsl

import (
	"fmt"
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"github.com/goadesign/goa/design"
)

// Attribute implements the attribute definition DSL. An attribute describes a data structure
// recursively. Attributes are used for describing request headers, parameters and payloads -
// response bodies and headers - media types and types. An attribute definition is recursive:
// attributes may include other attributes. At the basic level an attribute has a name,
// a type and optionally a default value and validation rules. The type of an attribute can be one of:
//
// * The primitive types Boolean, Integer, Number or String.
//
// * A type defined via the Type function.
//
// * A media type defined via the MediaType function.
//
// * An object described recursively with child attributes.
//
// * An array defined using the ArrayOf function.
//
// * An hashmap defined using the HashOf function.
//
// Attributes can be defined using the Attribute, Param, Member or Header functions depending
// on where the definition appears. The syntax for all these DSL is the same.
// Here are some examples:
//
//	Attribute("name")					// Defines an attribute of type String
//
//	Attribute("name", func() {
//		Pattern("^foo")					// Adds a validation rule to the attribute
//	})
//
//	Attribute("name", Integer)				// Defines an attribute of type Integer
//
//	Attribute("name", Integer, func() {
//		Default(42)					// With a default value
//	})
//
//	Attribute("name", Integer, "description")		// Specifies a description
//
//	Attribute("name", Integer, "description", func() {
//		Enum(1, 2)					// And validation rules
//	})
//
// Nested attributes:
//
//	Attribute("nested", func() {
//		Description("description")
//		Attribute("child")
//		Attribute("child2", func() {
//			// ....
//		})
//		Required("child")
//	})
//
// Here are all the valid usage of the Attribute function:
//
//	Attribute(name string, dataType DataType, description string, dsl func())
//
//	Attribute(name string, dataType DataType, description string)
//
//	Attribute(name string, dataType DataType, dsl func())
//
//	Attribute(name string, dataType DataType)
//
//	Attribute(name string, dsl func())	// dataType is String or Object (if DSL defines child attributes)
//
//	Attribute(name string)			// dataType is String
func Attribute(name string, args ...interface{}) {
	var parent *design.AttributeDefinition
	if at, ok := attributeDefinition(false); ok {
		parent = at
	} else if mt, ok := mediaTypeDefinition(true); ok {
		parent = mt.AttributeDefinition
	}

	if parent != nil {
		if parent.Type == nil {
			parent.Type = design.Object{}
		}
		if _, ok := parent.Type.(design.Object); !ok {
			ReportError("can't define child attributes on attribute of type %s", parent.Type.Name())
			return
		}

		var baseAttr *design.AttributeDefinition
		if parent.Reference != nil {
			if att, ok := parent.Reference.ToObject()[name]; ok {
				baseAttr = att.Dup()
			}
		}

		dataType, description, dsl := parseAttributeArgs(baseAttr, args...)
		if baseAttr != nil {
			if description != "" {
				baseAttr.Description = description
			}
			if dataType != nil {
				baseAttr.Type = dataType
			}
		} else {
			baseAttr = &design.AttributeDefinition{
				Type:        dataType,
				Description: description,
			}
		}
		baseAttr.Reference = parent.Reference
		if dsl != nil {
			ExecuteDSL(dsl, baseAttr)
		}
		if baseAttr.Type == nil {
			// DSL did not contain an "Attribute" declaration
			baseAttr.Type = design.String
		}
		parent.Type.(design.Object)[name] = baseAttr
	}
}

func parseAttributeArgs(baseAttr *design.AttributeDefinition, args ...interface{}) (design.DataType, string, func()) {
	var (
		dataType    design.DataType
		description string
		dsl         func()
		ok          bool
	)

	parseDataType := func(expected string, index int) {
		if dataType, ok = args[index].(design.DataType); !ok {
			invalidArgError(expected, args[index])
		}
	}
	parseDescription := func(expected string, index int) {
		if description, ok = args[index].(string); !ok {
			invalidArgError(expected, args[index])
		}
	}
	parseDSL := func(index int, success, failure func()) {
		if dsl, ok = args[index].(func()); ok {
			success()
		} else {
			failure()
		}
	}

	success := func() {}

	switch len(args) {
	case 0:
		if baseAttr != nil {
			dataType = baseAttr.Type
		} else {
			dataType = design.String
		}
	case 1:
		success = func() {
			if baseAttr != nil {
				dataType = baseAttr.Type
			}
		}
		parseDSL(0, success, func() { parseDataType("DataType or func()", 0) })
	case 2:
		parseDataType("DataType", 0)
		parseDSL(1, success, func() { parseDescription("string or func()", 1) })
	case 3:
		parseDataType("DataType", 0)
		parseDescription("string", 1)
		parseDSL(2, success, func() { invalidArgError("func()", args[2]) })
	default:
		ReportError("too many arguments in call to Attribute")
	}

	return dataType, description, dsl
}

// Header is an alias of Attribute.
func Header(name string, args ...interface{}) {
	Attribute(name, args...)
}

// Member is an alias of Attribute.
func Member(name string, args ...interface{}) {
	Attribute(name, args...)
}

// Param is an alias of Attribute.
func Param(name string, args ...interface{}) {
	Attribute(name, args...)
}

// Default sets the default value for an attribute.
func Default(def interface{}) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && !a.Type.IsCompatible(def) {
			ReportError("default value %#v is incompatible with attribute of type %s",
				def, a.Type.Name())
		} else {
			a.DefaultValue = def
		}
	}
}

// Enum adds a "enum" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor76.
func Enum(val ...interface{}) {
	if a, ok := attributeDefinition(true); ok {
		ok := true
		for i, v := range val {
			// When can a.Type be nil? glad you asked
			// There are two ways to write an Attribute declaration with the DSL that
			// don't set the type: with one argument - just the name - in which case the type
			// is set to String or with two arguments - the name and DSL. In this latter form
			// the type can end up being either String - if the DSL does not define any
			// attribute - or object if it does.
			// Why allowing this? because it's not always possible to specify the type of an
			// object - an object may just be declared inline to represent a substructure.
			// OK then why not assuming object and not allowing for string? because the DSL
			// where there's only one argument and the type is string implicitely is very
			// useful and common, for example to list attributes that refer to other attributes
			// such as responses that refer to responses defined at the API level or links that
			// refer to the media type attributes. So if the form that takes a DSL always ended
			// up defining an object we'd have a weird situation where one arg is string and
			// two args is object. Breaks the least surprise principle. Soooo long story
			// short the lesser evil seems to be to allow the ambiguity. Also tests like the
			// one below are really a convenience to the user and not a fundamental feature
			// - not checking in the case the type is not known yet is OK.
			if a.Type != nil && !a.Type.IsCompatible(v) {
				ReportError("value %#v at index #d is incompatible with attribute of type %s",
					v, i, a.Type.Name())
				ok = false
			}
		}
		if ok {
			a.Validations = append(a.Validations, &design.EnumValidationDefinition{Values: val})
		}
	}
}

// SupportedValidationFormats lists the supported formats for use with the
// Format DSL.
var SupportedValidationFormats = []string{
	"cidr",
	"date-time",
	"email",
	"hostname",
	"ipv4",
	"ipv6",
	"mac",
	"regexp",
	"uri",
}

// Format adds a "format" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor104.
// The formats supported by goa are:
//
// "date-time": RFC3339 date time
//
// "email": RFC5322 email address
//
// "hostname": RFC1035 internet host name
//
// "ipv4" and "ipv6": RFC2373 IPv4 and IPv6 address
//
// "uri": RFC3986 URI
//
// "mac": IEEE 802 MAC-48, EUI-48 or EUI-64 MAC address
//
// "cidr": RFC4632 or RFC4291 CIDR notation IP address
//
// "regexp": RE2 regular expression
func Format(f string) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && (a.Type.Kind() != design.StringKind && a.Type.Kind() != design.DateKind) {
			fmt.Println("skipping")
			incompatibleAttributeType("format", a.Type.Name(), "a string or date")
		} else {
			fmt.Println("Parsing")
			supported := false
			for _, s := range SupportedValidationFormats {
				if s == f {
					supported = true
					break
				}
			}
			if !supported {
				ReportError("unsupported format %#v, supported formats are: %s",
					f, strings.Join(SupportedValidationFormats, ", "))
			} else {
				a.Validations = append(a.Validations, &design.FormatValidationDefinition{Format: f})
			}
		}
	}
}

// Pattern adds a "pattern" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor33.
func Pattern(p string) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.StringKind {
			incompatibleAttributeType("pattern", a.Type.Name(), "a string")
		} else {
			_, err := regexp.Compile(p)
			if err != nil {
				ReportError("invalid pattern %#v, %s", p, err)
			} else {
				a.Validations = append(a.Validations, &design.PatternValidationDefinition{Pattern: p})
			}
		}
	}
}

// Minimum adds a "minimum" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor21.
func Minimum(val interface{}) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.IntegerKind && a.Type.Kind() != design.NumberKind {
			incompatibleAttributeType("minimum", a.Type.Name(), "an integer or a number")
		} else {
			var f float64
			switch v := val.(type) {
			case float32, float64, int, int8, int16, int32, int64, uint8, uint16, uint32, uint64:
				f = reflect.ValueOf(v).Convert(reflect.TypeOf(float64(0.0))).Float()
			case string:
				var err error
				f, err = strconv.ParseFloat(v, 64)
				if err != nil {
					ReportError("invalid number value %#v", v)
					return
				}
			default:
				ReportError("invalid number value %#v", v)
				return
			}
			a.Validations = append(a.Validations, &design.MinimumValidationDefinition{Min: f})
		}
	}
}

// Maximum adds a "maximum" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor17.
func Maximum(val interface{}) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.IntegerKind && a.Type.Kind() != design.NumberKind {
			incompatibleAttributeType("maximum", a.Type.Name(), "an integer or a number")
		} else {
			var f float64
			switch v := val.(type) {
			case float32, float64, int, int8, int16, int32, int64, uint8, uint16, uint32, uint64:
				f = reflect.ValueOf(v).Convert(reflect.TypeOf(float64(0.0))).Float()
			case string:
				var err error
				f, err = strconv.ParseFloat(v, 64)
				if err != nil {
					ReportError("invalid number value %#v", v)
					return
				}
			default:
				ReportError("invalid number value %#v", v)
				return
			}
			a.Validations = append(a.Validations, &design.MaximumValidationDefinition{Max: f})
		}
	}
}

// MinLength adss a "minItems" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor45.
func MinLength(val int) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.StringKind && a.Type.Kind() != design.ArrayKind {
			incompatibleAttributeType("minimum length", a.Type.Name(), "a string or an array")
		} else {
			a.Validations = append(a.Validations, &design.MinLengthValidationDefinition{MinLength: val})
		}
	}
}

// MaxLength adss a "maxItems" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor42.
func MaxLength(val int) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.StringKind && a.Type.Kind() != design.ArrayKind {
			incompatibleAttributeType("maximum length", a.Type.Name(), "a string or an array")
		} else {
			a.Validations = append(a.Validations, &design.MaxLengthValidationDefinition{MaxLength: val})
		}
	}
}

// Required adds a "required" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor61.
func Required(names ...string) {
	var at *design.AttributeDefinition
	if a, ok := attributeDefinition(false); ok {
		at = a
	} else if mt, ok := mediaTypeDefinition(true); ok {
		at = mt.AttributeDefinition
	} else {
		return
	}
	if at.Type != nil && at.Type.Kind() != design.ObjectKind {
		incompatibleAttributeType("required", at.Type.Name(), "an object")
	} else {
		at.Validations = append(at.Validations, &design.RequiredValidationDefinition{Names: names})
	}
}

// incompatibleAttributeType reports an error for validations defined on
// incompatible attributes (e.g. max value on string).
func incompatibleAttributeType(validation, actual, expected string) {
	ReportError("invalid %s validation definition: attribute must be %s (but type is %s)",
		validation, expected, actual)
}
<commit_msg>remove debug print<commit_after>package dsl

import (
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"github.com/goadesign/goa/design"
)

// Attribute implements the attribute definition DSL. An attribute describes a data structure
// recursively. Attributes are used for describing request headers, parameters and payloads -
// response bodies and headers - media types and types. An attribute definition is recursive:
// attributes may include other attributes. At the basic level an attribute has a name,
// a type and optionally a default value and validation rules. The type of an attribute can be one of:
//
// * The primitive types Boolean, Integer, Number or String.
//
// * A type defined via the Type function.
//
// * A media type defined via the MediaType function.
//
// * An object described recursively with child attributes.
//
// * An array defined using the ArrayOf function.
//
// * An hashmap defined using the HashOf function.
//
// Attributes can be defined using the Attribute, Param, Member or Header functions depending
// on where the definition appears. The syntax for all these DSL is the same.
// Here are some examples:
//
//	Attribute("name")					// Defines an attribute of type String
//
//	Attribute("name", func() {
//		Pattern("^foo")					// Adds a validation rule to the attribute
//	})
//
//	Attribute("name", Integer)				// Defines an attribute of type Integer
//
//	Attribute("name", Integer, func() {
//		Default(42)					// With a default value
//	})
//
//	Attribute("name", Integer, "description")		// Specifies a description
//
//	Attribute("name", Integer, "description", func() {
//		Enum(1, 2)					// And validation rules
//	})
//
// Nested attributes:
//
//	Attribute("nested", func() {
//		Description("description")
//		Attribute("child")
//		Attribute("child2", func() {
//			// ....
//		})
//		Required("child")
//	})
//
// Here are all the valid usage of the Attribute function:
//
//	Attribute(name string, dataType DataType, description string, dsl func())
//
//	Attribute(name string, dataType DataType, description string)
//
//	Attribute(name string, dataType DataType, dsl func())
//
//	Attribute(name string, dataType DataType)
//
//	Attribute(name string, dsl func())	// dataType is String or Object (if DSL defines child attributes)
//
//	Attribute(name string)			// dataType is String
func Attribute(name string, args ...interface{}) {
	var parent *design.AttributeDefinition
	if at, ok := attributeDefinition(false); ok {
		parent = at
	} else if mt, ok := mediaTypeDefinition(true); ok {
		parent = mt.AttributeDefinition
	}

	if parent != nil {
		if parent.Type == nil {
			parent.Type = design.Object{}
		}
		if _, ok := parent.Type.(design.Object); !ok {
			ReportError("can't define child attributes on attribute of type %s", parent.Type.Name())
			return
		}

		var baseAttr *design.AttributeDefinition
		if parent.Reference != nil {
			if att, ok := parent.Reference.ToObject()[name]; ok {
				baseAttr = att.Dup()
			}
		}

		dataType, description, dsl := parseAttributeArgs(baseAttr, args...)
		if baseAttr != nil {
			if description != "" {
				baseAttr.Description = description
			}
			if dataType != nil {
				baseAttr.Type = dataType
			}
		} else {
			baseAttr = &design.AttributeDefinition{
				Type:        dataType,
				Description: description,
			}
		}
		baseAttr.Reference = parent.Reference
		if dsl != nil {
			ExecuteDSL(dsl, baseAttr)
		}
		if baseAttr.Type == nil {
			// DSL did not contain an "Attribute" declaration
			baseAttr.Type = design.String
		}
		parent.Type.(design.Object)[name] = baseAttr
	}
}

func parseAttributeArgs(baseAttr *design.AttributeDefinition, args ...interface{}) (design.DataType, string, func()) {
	var (
		dataType    design.DataType
		description string
		dsl         func()
		ok          bool
	)

	parseDataType := func(expected string, index int) {
		if dataType, ok = args[index].(design.DataType); !ok {
			invalidArgError(expected, args[index])
		}
	}
	parseDescription := func(expected string, index int) {
		if description, ok = args[index].(string); !ok {
			invalidArgError(expected, args[index])
		}
	}
	parseDSL := func(index int, success, failure func()) {
		if dsl, ok = args[index].(func()); ok {
			success()
		} else {
			failure()
		}
	}

	success := func() {}

	switch len(args) {
	case 0:
		if baseAttr != nil {
			dataType = baseAttr.Type
		} else {
			dataType = design.String
		}
	case 1:
		success = func() {
			if baseAttr != nil {
				dataType = baseAttr.Type
			}
		}
		parseDSL(0, success, func() { parseDataType("DataType or func()", 0) })
	case 2:
		parseDataType("DataType", 0)
		parseDSL(1, success, func() { parseDescription("string or func()", 1) })
	case 3:
		parseDataType("DataType", 0)
		parseDescription("string", 1)
		parseDSL(2, success, func() { invalidArgError("func()", args[2]) })
	default:
		ReportError("too many arguments in call to Attribute")
	}

	return dataType, description, dsl
}

// Header is an alias of Attribute.
func Header(name string, args ...interface{}) {
	Attribute(name, args...)
}

// Member is an alias of Attribute.
func Member(name string, args ...interface{}) {
	Attribute(name, args...)
}

// Param is an alias of Attribute.
func Param(name string, args ...interface{}) {
	Attribute(name, args...)
}

// Default sets the default value for an attribute.
func Default(def interface{}) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && !a.Type.IsCompatible(def) {
			ReportError("default value %#v is incompatible with attribute of type %s",
				def, a.Type.Name())
		} else {
			a.DefaultValue = def
		}
	}
}

// Enum adds a "enum" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor76.
func Enum(val ...interface{}) {
	if a, ok := attributeDefinition(true); ok {
		ok := true
		for i, v := range val {
			// When can a.Type be nil? glad you asked
			// There are two ways to write an Attribute declaration with the DSL that
			// don't set the type: with one argument - just the name - in which case the type
			// is set to String or with two arguments - the name and DSL. In this latter form
			// the type can end up being either String - if the DSL does not define any
			// attribute - or object if it does.
			// Why allowing this? because it's not always possible to specify the type of an
			// object - an object may just be declared inline to represent a substructure.
			// OK then why not assuming object and not allowing for string? because the DSL
			// where there's only one argument and the type is string implicitely is very
			// useful and common, for example to list attributes that refer to other attributes
			// such as responses that refer to responses defined at the API level or links that
			// refer to the media type attributes. So if the form that takes a DSL always ended
			// up defining an object we'd have a weird situation where one arg is string and
			// two args is object. Breaks the least surprise principle. Soooo long story
			// short the lesser evil seems to be to allow the ambiguity. Also tests like the
			// one below are really a convenience to the user and not a fundamental feature
			// - not checking in the case the type is not known yet is OK.
			if a.Type != nil && !a.Type.IsCompatible(v) {
				ReportError("value %#v at index #d is incompatible with attribute of type %s",
					v, i, a.Type.Name())
				ok = false
			}
		}
		if ok {
			a.Validations = append(a.Validations, &design.EnumValidationDefinition{Values: val})
		}
	}
}

// SupportedValidationFormats lists the supported formats for use with the
// Format DSL.
var SupportedValidationFormats = []string{
	"cidr",
	"date-time",
	"email",
	"hostname",
	"ipv4",
	"ipv6",
	"mac",
	"regexp",
	"uri",
}

// Format adds a "format" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor104.
// The formats supported by goa are:
//
// "date-time": RFC3339 date time
//
// "email": RFC5322 email address
//
// "hostname": RFC1035 internet host name
//
// "ipv4" and "ipv6": RFC2373 IPv4 and IPv6 address
//
// "uri": RFC3986 URI
//
// "mac": IEEE 802 MAC-48, EUI-48 or EUI-64 MAC address
//
// "cidr": RFC4632 or RFC4291 CIDR notation IP address
//
// "regexp": RE2 regular expression
func Format(f string) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && (a.Type.Kind() != design.StringKind && a.Type.Kind() != design.DateKind) {
			incompatibleAttributeType("format", a.Type.Name(), "a string or date")
		} else {
			supported := false
			for _, s := range SupportedValidationFormats {
				if s == f {
					supported = true
					break
				}
			}
			if !supported {
				ReportError("unsupported format %#v, supported formats are: %s",
					f, strings.Join(SupportedValidationFormats, ", "))
			} else {
				a.Validations = append(a.Validations, &design.FormatValidationDefinition{Format: f})
			}
		}
	}
}

// Pattern adds a "pattern" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor33.
func Pattern(p string) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.StringKind {
			incompatibleAttributeType("pattern", a.Type.Name(), "a string")
		} else {
			_, err := regexp.Compile(p)
			if err != nil {
				ReportError("invalid pattern %#v, %s", p, err)
			} else {
				a.Validations = append(a.Validations, &design.PatternValidationDefinition{Pattern: p})
			}
		}
	}
}

// Minimum adds a "minimum" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor21.
func Minimum(val interface{}) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.IntegerKind && a.Type.Kind() != design.NumberKind {
			incompatibleAttributeType("minimum", a.Type.Name(), "an integer or a number")
		} else {
			var f float64
			switch v := val.(type) {
			case float32, float64, int, int8, int16, int32, int64, uint8, uint16, uint32, uint64:
				f = reflect.ValueOf(v).Convert(reflect.TypeOf(float64(0.0))).Float()
			case string:
				var err error
				f, err = strconv.ParseFloat(v, 64)
				if err != nil {
					ReportError("invalid number value %#v", v)
					return
				}
			default:
				ReportError("invalid number value %#v", v)
				return
			}
			a.Validations = append(a.Validations, &design.MinimumValidationDefinition{Min: f})
		}
	}
}

// Maximum adds a "maximum" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor17.
func Maximum(val interface{}) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.IntegerKind && a.Type.Kind() != design.NumberKind {
			incompatibleAttributeType("maximum", a.Type.Name(), "an integer or a number")
		} else {
			var f float64
			switch v := val.(type) {
			case float32, float64, int, int8, int16, int32, int64, uint8, uint16, uint32, uint64:
				f = reflect.ValueOf(v).Convert(reflect.TypeOf(float64(0.0))).Float()
			case string:
				var err error
				f, err = strconv.ParseFloat(v, 64)
				if err != nil {
					ReportError("invalid number value %#v", v)
					return
				}
			default:
				ReportError("invalid number value %#v", v)
				return
			}
			a.Validations = append(a.Validations, &design.MaximumValidationDefinition{Max: f})
		}
	}
}

// MinLength adss a "minItems" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor45.
func MinLength(val int) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.StringKind && a.Type.Kind() != design.ArrayKind {
			incompatibleAttributeType("minimum length", a.Type.Name(), "a string or an array")
		} else {
			a.Validations = append(a.Validations, &design.MinLengthValidationDefinition{MinLength: val})
		}
	}
}

// MaxLength adss a "maxItems" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor42.
func MaxLength(val int) {
	if a, ok := attributeDefinition(true); ok {
		if a.Type != nil && a.Type.Kind() != design.StringKind && a.Type.Kind() != design.ArrayKind {
			incompatibleAttributeType("maximum length", a.Type.Name(), "a string or an array")
		} else {
			a.Validations = append(a.Validations, &design.MaxLengthValidationDefinition{MaxLength: val})
		}
	}
}

// Required adds a "required" validation to the attribute.
// See http://json-schema.org/latest/json-schema-validation.html#anchor61.
func Required(names ...string) {
	var at *design.AttributeDefinition
	if a, ok := attributeDefinition(false); ok {
		at = a
	} else if mt, ok := mediaTypeDefinition(true); ok {
		at = mt.AttributeDefinition
	} else {
		return
	}
	if at.Type != nil && at.Type.Kind() != design.ObjectKind {
		incompatibleAttributeType("required", at.Type.Name(), "an object")
	} else {
		at.Validations = append(at.Validations, &design.RequiredValidationDefinition{Names: names})
	}
}

// incompatibleAttributeType reports an error for validations defined on
// incompatible attributes (e.g. max value on string).
func incompatibleAttributeType(validation, actual, expected string) {
	ReportError("invalid %s validation definition: attribute must be %s (but type is %s)",
		validation, expected, actual)
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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
//
// 	https://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 headers

import (
	"context"
	"net/http"
	"testing"

	"github.com/google/go-cmp/cmp"

	"github.com/google/go-safeweb/internal/requesttesting"
)

func TestHeaderParsing(t *testing.T) {
	var tests = []struct {
		name    string
		request []byte
		want    map[string][]string
	}{
		{
			name: "Basic",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A: B\r\n" +
				"\r\n"),
			want: map[string][]string{"A": []string{"B"}},
		},
		{
			name: "Ordering",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A: X\r\n" +
				"A: Y\r\n" +
				"\r\n"),
			want: map[string][]string{"A": []string{"X", "Y"}},
		},
		{
			name: "Casing",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"BaBaBaBa-BaBaBaBa-BaBaBa: xXxXxXxX\r\n" +
				"cDcDcDcD-cDcDcDcD-cDcDcD: YyYyYyYy\r\n" +
				"\r\n"),
			want: map[string][]string{
				"Babababa-Babababa-Bababa": []string{"xXxXxXxX"},
				"Cdcdcdcd-Cdcdcdcd-Cdcdcd": []string{"YyYyYyYy"},
			},
		},
		{
			name: "CasingOrdering",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"a: X\r\n" +
				"A: Y\r\n" +
				"a: Z\r\n" +
				"A: W\r\n" +
				"\r\n"),
			want: map[string][]string{"A": []string{"X", "Y", "Z", "W"}},
		},
		{
			name: "MultiLineHeaderSpace",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AAAA: aaaa aaa\r\n" +
				" aaa aaa\r\n" +
				"\r\n"),
			want: http.Header{"Aaaa": []string{"aaaa aaa aaa aaa"}},
		},
		{
			name: "MultiLineHeaderTab",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AAAA: aaaa aaa\r\n" +
				"\taaa aaa\r\n" +
				"\r\n"),
			want: http.Header{"Aaaa": []string{"aaaa aaa aaa aaa"}},
		},
		{
			name: "MultiLineHeaderManyLines",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AAAA: aaaa aaa\r\n" +
				" aaa\r\n" +
				" aaa\r\n" +
				"\r\n"),
			want: http.Header{"Aaaa": []string{"aaaa aaa aaa aaa"}},
		},
		{
			name: "UnicodeValue",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A: \xf0\x9f\xa5\xb3\r\n" +
				"\r\n"),
			want: http.Header{"A": []string{"\xf0\x9f\xa5\xb3"}},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			resp, err := requesttesting.MakeRequest(context.Background(), tt.request, func(r *http.Request) {
				if diff := cmp.Diff(tt.want, map[string][]string(r.Header)); diff != "" {
					t.Errorf("r.Header mismatch (-want +got):\n%s", diff)
				}
			})
			if err != nil {
				t.Fatalf("MakeRequest() got err: %v", err)
			}

			if got, want := extractStatus(resp), statusOK; got != want {
				t.Errorf("status code got: %q want: %q", got, want)
			}
		})
	}
}

func TestStatusCode(t *testing.T) {
	var tests = []struct {
		name    string
		request []byte
		want    string
	}{
		{
			name: "MultilineContinuationOnFirstLine",
			request: []byte("GET / HTTP/1.1\r\n" +
				" A: a\r\n" +
				"Host: localhost:8080\r\n" +
				"\r\n"),
			want: statusBadRequest,
		},
		{
			name: "MultilineHeaderName",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AA\r\n" +
				" AA: aaaa\r\n" +
				"\r\n"),
			want: statusBadRequest,
		},
		{
			name: "MultilineHeaderBeforeColon",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A\r\n" +
				" : a\r\n" +
				"\r\n"),
			want: statusBadRequest,
		},
		{
			name: "UnicodeHeaderName",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"\xf0\x9f\xa5\xb3: a\r\n" +
				"\r\n"),
			want: statusInvalidHeaderName,
		},
		{
			name: "SpecialCharactersHeaderName",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"&%*(#@%()): a\r\n" +
				"\r\n"),
			want: statusInvalidHeaderName,
		},
		{
			name: "NullByteInHeaderName",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A\x00A: a\r\n" +
				"\r\n"),
			want: statusInvalidHeaderName,
		},
		{
			name: "NullByteInHeaderValue",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AA: a\x00a\r\n" +
				"\r\n"),
			want: statusInvalidHeaderValue,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			resp, err := requesttesting.MakeRequest(context.Background(), tt.request, nil)
			if err != nil {
				t.Fatalf("MakeRequest() got err: %v", err)
			}

			if got := extractStatus(resp); got != tt.want {
				t.Errorf("status code got: %q want: %q", got, tt.want)
			}
		})
	}
}
<commit_msg>Added test for ordering of http.Header.Values()<commit_after>// Copyright 2020 Google LLC
//
// 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
//
// 	https://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 headers

import (
	"context"
	"net/http"
	"testing"

	"github.com/google/go-cmp/cmp"

	"github.com/google/go-safeweb/internal/requesttesting"
)

func TestHeaderParsing(t *testing.T) {
	var tests = []struct {
		name    string
		request []byte
		want    map[string][]string
	}{
		{
			name: "Basic",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A: B\r\n" +
				"\r\n"),
			want: map[string][]string{"A": []string{"B"}},
		},
		{
			name: "Ordering",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A: X\r\n" +
				"A: Y\r\n" +
				"\r\n"),
			want: map[string][]string{"A": []string{"X", "Y"}},
		},
		{
			name: "Casing",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"BaBaBaBa-BaBaBaBa-BaBaBa: xXxXxXxX\r\n" +
				"cDcDcDcD-cDcDcDcD-cDcDcD: YyYyYyYy\r\n" +
				"\r\n"),
			want: map[string][]string{
				"Babababa-Babababa-Bababa": []string{"xXxXxXxX"},
				"Cdcdcdcd-Cdcdcdcd-Cdcdcd": []string{"YyYyYyYy"},
			},
		},
		{
			name: "CasingOrdering",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"a: X\r\n" +
				"A: Y\r\n" +
				"a: Z\r\n" +
				"A: W\r\n" +
				"\r\n"),
			want: map[string][]string{"A": []string{"X", "Y", "Z", "W"}},
		},
		{
			name: "MultiLineHeaderSpace",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AAAA: aaaa aaa\r\n" +
				" aaa aaa\r\n" +
				"\r\n"),
			want: http.Header{"Aaaa": []string{"aaaa aaa aaa aaa"}},
		},
		{
			name: "MultiLineHeaderTab",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AAAA: aaaa aaa\r\n" +
				"\taaa aaa\r\n" +
				"\r\n"),
			want: http.Header{"Aaaa": []string{"aaaa aaa aaa aaa"}},
		},
		{
			name: "MultiLineHeaderManyLines",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AAAA: aaaa aaa\r\n" +
				" aaa\r\n" +
				" aaa\r\n" +
				"\r\n"),
			want: http.Header{"Aaaa": []string{"aaaa aaa aaa aaa"}},
		},
		{
			name: "UnicodeValue",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A: \xf0\x9f\xa5\xb3\r\n" +
				"\r\n"),
			want: http.Header{"A": []string{"\xf0\x9f\xa5\xb3"}},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			resp, err := requesttesting.MakeRequest(context.Background(), tt.request, func(r *http.Request) {
				if diff := cmp.Diff(tt.want, map[string][]string(r.Header)); diff != "" {
					t.Errorf("r.Header mismatch (-want +got):\n%s", diff)
				}
			})
			if err != nil {
				t.Fatalf("MakeRequest() got err: %v", err)
			}

			if got, want := extractStatus(resp), statusOK; got != want {
				t.Errorf("status code got: %q want: %q", got, want)
			}
		})
	}
}

func TestStatusCode(t *testing.T) {
	var tests = []struct {
		name    string
		request []byte
		want    string
	}{
		{
			name: "MultilineContinuationOnFirstLine",
			request: []byte("GET / HTTP/1.1\r\n" +
				" A: a\r\n" +
				"Host: localhost:8080\r\n" +
				"\r\n"),
			want: statusBadRequest,
		},
		{
			name: "MultilineHeaderName",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AA\r\n" +
				" AA: aaaa\r\n" +
				"\r\n"),
			want: statusBadRequest,
		},
		{
			name: "MultilineHeaderBeforeColon",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A\r\n" +
				" : a\r\n" +
				"\r\n"),
			want: statusBadRequest,
		},
		{
			name: "UnicodeHeaderName",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"\xf0\x9f\xa5\xb3: a\r\n" +
				"\r\n"),
			want: statusInvalidHeaderName,
		},
		{
			name: "SpecialCharactersHeaderName",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"&%*(#@%()): a\r\n" +
				"\r\n"),
			want: statusInvalidHeaderName,
		},
		{
			name: "NullByteInHeaderName",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A\x00A: a\r\n" +
				"\r\n"),
			want: statusInvalidHeaderName,
		},
		{
			name: "NullByteInHeaderValue",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"AA: a\x00a\r\n" +
				"\r\n"),
			want: statusInvalidHeaderValue,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			resp, err := requesttesting.MakeRequest(context.Background(), tt.request, nil)
			if err != nil {
				t.Fatalf("MakeRequest() got err: %v", err)
			}

			if got := extractStatus(resp); got != tt.want {
				t.Errorf("status code got: %q want: %q", got, tt.want)
			}
		})
	}
}

func TestValues(t *testing.T) {
	var tests = []struct {
		name    string
		request []byte
		want    []string
	}{
		{
			name: "Ordering1",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A: X\r\n" +
				"A: Y\r\n" +
				"\r\n"),
			want: []string{"X", "Y"},
		},
		{
			name: "Ordering2",
			request: []byte("GET / HTTP/1.1\r\n" +
				"Host: localhost:8080\r\n" +
				"A: Y\r\n" +
				"A: X\r\n" +
				"\r\n"),
			want: []string{"Y", "X"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			resp, err := requesttesting.MakeRequest(context.Background(), tt.request, func(r *http.Request) {
				if diff := cmp.Diff(tt.want, r.Header.Values("A")); diff != "" {
					t.Errorf(`r.Header.Values("A") mismatch (-want +got):\n%s`, diff)
				}
			})
			if err != nil {
				t.Fatalf("MakeRequest() got err: %v want: nil", err)
			}

			if got, want := extractStatus(resp), statusOK; got != want {
				t.Errorf("status code got: %q want: %q", got, want)
			}
		})
	}
}
<|endoftext|>


<commit_before>package v1datatransfer

import (
	"bytes"
	"context"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"path"
	"time"

	v1batch "github.com/nerdalize/nerd/nerd/client/batch/v1"
	v1payload "github.com/nerdalize/nerd/nerd/client/batch/v1/payload"
	v1data "github.com/nerdalize/nerd/nerd/service/datatransfer/v1/client"
	v1datapayload "github.com/nerdalize/nerd/nerd/service/datatransfer/v1/client/payload"
	"github.com/pkg/errors"
	"github.com/restic/chunker"
)

type uploadProcess struct {
	batchClient       v1batch.ClientUploadInterface
	dataClient        *v1data.Client
	dataset           v1payload.DatasetSummary
	heartbeatInterval time.Duration
	localDir          string
	concurrency       int
	progressCh        chan<- int64
}

//start starts the upload process
func (p *uploadProcess) start(ctx context.Context) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	type countResult struct {
		total int64
		err   error
	}

	// pipeline: tar | count | chunk+upload | index
	tarCountPipe := newPipe()
	defer tarCountPipe.w.Close()
	countChunksPipe := newPipe()
	chunksIndexPipe := newPipe()
	countReader := io.TeeReader(tarCountPipe.r, countChunksPipe.w)

	doneCh := make(chan error, 4)
	countCh := make(chan countResult, 1)

	// heartbeat
	go sendHeartbeat(ctx, cancel, doneCh, p.batchClient, p.dataset.ProjectID, p.dataset.DatasetID, p.heartbeatInterval)

	// chunks | index
	go func() {
		doneCh <- uploadIndex(ctx, p.dataClient, chunksIndexPipe.r, p.dataset.Bucket, p.dataset.DatasetRoot)
	}()
	// count | chunks
	go func() {
		defer chunksIndexPipe.w.Close()
		kw := v1data.NewIndexWriter(chunksIndexPipe.w)
		err := uploadChunks(ctx, p.dataClient, countChunksPipe.r, kw, p.dataset.Bucket, p.dataset.ProjectRoot, p.concurrency, p.progressCh)
		if err != nil {
			doneCh <- err
		}
	}()
	// tar | count
	go func() {
		total, err := countBytes(ctx, countReader)
		countCh <- countResult{
			err:   err,
			total: total,
		}
	}()
	// tar
	go func() {
		defer tarCountPipe.w.Close()
		defer countChunksPipe.w.Close()
		err := tardir(ctx, p.localDir, tarCountPipe.w)
		if err != nil {
			doneCh <- err
		}
	}()

	err := <-doneCh
	if err != nil {
		return err
	}

	cres := <-countCh
	if cres.err != nil {
		return errors.Wrap(err, "failed to calculate dataset size")
	}
	err = uploadMetadata(ctx, p.dataClient, cres.total, p.dataset.Bucket, p.dataset.DatasetRoot)
	if err != nil {
		return err
	}

	_, err = p.batchClient.SendUploadSuccess(p.dataset.ProjectID, p.dataset.DatasetID)
	if err != nil {
		return errors.Wrap(err, "failed to send dataset success message")
	}

	return nil
}

//uploadChunks uploads data from r in a chunked way
func uploadChunks(ctx context.Context, dataClient *v1data.Client, r io.Reader, kw v1data.KeyWriter, bucket, root string, concurrency int, progressCh chan<- int64) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	chkr := chunker.New(r, chunker.Pol(v1data.UploadPolynomal))
	type result struct {
		err error
		k   v1data.Key
	}

	type item struct {
		chunk []byte
		size  int64
		resCh chan *result
		err   error
	}

	work := func(it *item) {
		k := v1data.Key(sha256.Sum256(it.chunk)) //hash
		key := path.Join(root, k.ToString())
		exists, err := dataClient.Exists(ctx, bucket, key) //check existence
		if err != nil {
			select {
			case <-ctx.Done():
			case it.resCh <- &result{errors.Wrapf(err, "failed to check existence of '%x'", k), v1data.ZeroKey}:
			}
			return
		}

		if !exists {
			err = dataClient.Upload(ctx, bucket, key, bytes.NewReader(it.chunk)) //if not exists put
			if err != nil {
				select {
				case <-ctx.Done():
				case it.resCh <- &result{errors.Wrapf(err, "failed to put chunk '%x'", k), v1data.ZeroKey}:
				}
				return
			}
		}
		if progressCh != nil {
			select {
			case <-ctx.Done():
			case progressCh <- int64(len(it.chunk)):
			}
		}

		it.resCh <- &result{nil, k}
	}

	//fan out
	itemCh := make(chan *item, concurrency)
	go func() {
		defer close(itemCh)
		buf := make([]byte, chunker.MaxSize)
		for {
			chunk, err := chkr.Next(buf)
			if err != nil {
				if err != io.EOF {
					select {
					case <-ctx.Done():
					case itemCh <- &item{err: err}:
					}
				}
				return
			}

			it := &item{
				chunk: make([]byte, chunk.Length),
				resCh: make(chan *result),
			}

			copy(it.chunk, chunk.Data) //underlying buffer is switched out

			select {
			case <-ctx.Done():
			case itemCh <- it: //send to fan-in thread for syncing results
				go work(it) //create work
			}
		}
	}()

	//fan-in
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case it := <-itemCh:
			if it == nil {
				return nil
			}
			if it.err != nil {
				return errors.Wrapf(it.err, "failed to iterate")
			}

			select {
			case <-ctx.Done():
				return ctx.Err()
			case res := <-it.resCh:
				if res.err != nil {
					return res.err
				}

				err := kw.WriteKey(res.k)
				if err != nil {
					return errors.Wrapf(err, "failed to write key")
				}
			}
		}
	}

	return nil
}

//uploadIndex uploads the index object with all the keys
func uploadIndex(ctx context.Context, dataClient *v1data.Client, r io.Reader, bucket, root string) error {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return errors.Wrap(err, "failed to read keys")
	}
	return dataClient.Upload(ctx, bucket, path.Join(root, v1data.IndexObjectKey), bytes.NewReader(b))
}

//uploadMetadata uploads the metadata object
func uploadMetadata(ctx context.Context, dataClient *v1data.Client, total int64, bucket, root string) error {
	metadata := &v1datapayload.Metadata{
		Size:    total,
		Created: time.Now(),
		Updated: time.Now(),
	}
	err := dataClient.MetadataUpload(ctx, bucket, root, metadata)
	if err != nil {
		return errors.Wrap(err, "failed to upload metadata")
	}
	return nil
}

//sendHeartbeat sends a heartbeat and sleeps for the given interval
func sendHeartbeat(ctx context.Context, cancel context.CancelFunc, doneCh chan error, batchClient v1batch.ClientUploadInterface, projectID, datasetID string, interval time.Duration) {
	defer fmt.Println("heartbeat done")
	ticker := time.Tick(interval)
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker:
			out, err := batchClient.SendUploadHeartbeat(projectID, datasetID)
			if err == nil && out.HasExpired {
				doneCh <- fmt.Errorf("upload failed because the server could not be reached for too long")
				cancel()
			}
		}
	}
}
<commit_msg>Remove print statement<commit_after>package v1datatransfer

import (
	"bytes"
	"context"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"path"
	"time"

	v1batch "github.com/nerdalize/nerd/nerd/client/batch/v1"
	v1payload "github.com/nerdalize/nerd/nerd/client/batch/v1/payload"
	v1data "github.com/nerdalize/nerd/nerd/service/datatransfer/v1/client"
	v1datapayload "github.com/nerdalize/nerd/nerd/service/datatransfer/v1/client/payload"
	"github.com/pkg/errors"
	"github.com/restic/chunker"
)

type uploadProcess struct {
	batchClient       v1batch.ClientUploadInterface
	dataClient        *v1data.Client
	dataset           v1payload.DatasetSummary
	heartbeatInterval time.Duration
	localDir          string
	concurrency       int
	progressCh        chan<- int64
}

//start starts the upload process
func (p *uploadProcess) start(ctx context.Context) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	type countResult struct {
		total int64
		err   error
	}

	// pipeline: tar | count | chunk+upload | index
	tarCountPipe := newPipe()
	defer tarCountPipe.w.Close()
	countChunksPipe := newPipe()
	chunksIndexPipe := newPipe()
	countReader := io.TeeReader(tarCountPipe.r, countChunksPipe.w)

	doneCh := make(chan error, 4)
	countCh := make(chan countResult, 1)

	// heartbeat
	go sendHeartbeat(ctx, cancel, doneCh, p.batchClient, p.dataset.ProjectID, p.dataset.DatasetID, p.heartbeatInterval)

	// chunks | index
	go func() {
		doneCh <- uploadIndex(ctx, p.dataClient, chunksIndexPipe.r, p.dataset.Bucket, p.dataset.DatasetRoot)
	}()
	// count | chunks
	go func() {
		defer chunksIndexPipe.w.Close()
		kw := v1data.NewIndexWriter(chunksIndexPipe.w)
		err := uploadChunks(ctx, p.dataClient, countChunksPipe.r, kw, p.dataset.Bucket, p.dataset.ProjectRoot, p.concurrency, p.progressCh)
		if err != nil {
			doneCh <- err
		}
	}()
	// tar | count
	go func() {
		total, err := countBytes(ctx, countReader)
		countCh <- countResult{
			err:   err,
			total: total,
		}
	}()
	// tar
	go func() {
		defer tarCountPipe.w.Close()
		defer countChunksPipe.w.Close()
		err := tardir(ctx, p.localDir, tarCountPipe.w)
		if err != nil {
			doneCh <- err
		}
	}()

	err := <-doneCh
	if err != nil {
		return err
	}

	cres := <-countCh
	if cres.err != nil {
		return errors.Wrap(err, "failed to calculate dataset size")
	}
	err = uploadMetadata(ctx, p.dataClient, cres.total, p.dataset.Bucket, p.dataset.DatasetRoot)
	if err != nil {
		return err
	}

	_, err = p.batchClient.SendUploadSuccess(p.dataset.ProjectID, p.dataset.DatasetID)
	if err != nil {
		return errors.Wrap(err, "failed to send dataset success message")
	}

	return nil
}

//uploadChunks uploads data from r in a chunked way
func uploadChunks(ctx context.Context, dataClient *v1data.Client, r io.Reader, kw v1data.KeyWriter, bucket, root string, concurrency int, progressCh chan<- int64) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	chkr := chunker.New(r, chunker.Pol(v1data.UploadPolynomal))
	type result struct {
		err error
		k   v1data.Key
	}

	type item struct {
		chunk []byte
		size  int64
		resCh chan *result
		err   error
	}

	work := func(it *item) {
		k := v1data.Key(sha256.Sum256(it.chunk)) //hash
		key := path.Join(root, k.ToString())
		exists, err := dataClient.Exists(ctx, bucket, key) //check existence
		if err != nil {
			select {
			case <-ctx.Done():
			case it.resCh <- &result{errors.Wrapf(err, "failed to check existence of '%x'", k), v1data.ZeroKey}:
			}
			return
		}

		if !exists {
			err = dataClient.Upload(ctx, bucket, key, bytes.NewReader(it.chunk)) //if not exists put
			if err != nil {
				select {
				case <-ctx.Done():
				case it.resCh <- &result{errors.Wrapf(err, "failed to put chunk '%x'", k), v1data.ZeroKey}:
				}
				return
			}
		}
		if progressCh != nil {
			select {
			case <-ctx.Done():
			case progressCh <- int64(len(it.chunk)):
			}
		}

		it.resCh <- &result{nil, k}
	}

	//fan out
	itemCh := make(chan *item, concurrency)
	go func() {
		defer close(itemCh)
		buf := make([]byte, chunker.MaxSize)
		for {
			chunk, err := chkr.Next(buf)
			if err != nil {
				if err != io.EOF {
					select {
					case <-ctx.Done():
					case itemCh <- &item{err: err}:
					}
				}
				return
			}

			it := &item{
				chunk: make([]byte, chunk.Length),
				resCh: make(chan *result),
			}

			copy(it.chunk, chunk.Data) //underlying buffer is switched out

			select {
			case <-ctx.Done():
			case itemCh <- it: //send to fan-in thread for syncing results
				go work(it) //create work
			}
		}
	}()

	//fan-in
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case it := <-itemCh:
			if it == nil {
				return nil
			}
			if it.err != nil {
				return errors.Wrapf(it.err, "failed to iterate")
			}

			select {
			case <-ctx.Done():
				return ctx.Err()
			case res := <-it.resCh:
				if res.err != nil {
					return res.err
				}

				err := kw.WriteKey(res.k)
				if err != nil {
					return errors.Wrapf(err, "failed to write key")
				}
			}
		}
	}

	return nil
}

//uploadIndex uploads the index object with all the keys
func uploadIndex(ctx context.Context, dataClient *v1data.Client, r io.Reader, bucket, root string) error {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return errors.Wrap(err, "failed to read keys")
	}
	return dataClient.Upload(ctx, bucket, path.Join(root, v1data.IndexObjectKey), bytes.NewReader(b))
}

//uploadMetadata uploads the metadata object
func uploadMetadata(ctx context.Context, dataClient *v1data.Client, total int64, bucket, root string) error {
	metadata := &v1datapayload.Metadata{
		Size:    total,
		Created: time.Now(),
		Updated: time.Now(),
	}
	err := dataClient.MetadataUpload(ctx, bucket, root, metadata)
	if err != nil {
		return errors.Wrap(err, "failed to upload metadata")
	}
	return nil
}

//sendHeartbeat sends a heartbeat and sleeps for the given interval
func sendHeartbeat(ctx context.Context, cancel context.CancelFunc, doneCh chan error, batchClient v1batch.ClientUploadInterface, projectID, datasetID string, interval time.Duration) {
	ticker := time.Tick(interval)
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker:
			out, err := batchClient.SendUploadHeartbeat(projectID, datasetID)
			if err == nil && out.HasExpired {
				doneCh <- fmt.Errorf("upload failed because the server could not be reached for too long")
				cancel()
			}
		}
	}
}
<|endoftext|>