<commit_before>// Copyright 2015 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

package hamlib

import (
	"fmt"
	"io"
	"net"
	"net/textproto"
	"strconv"
	"strings"
	"sync"
)

const DefaultTCPAddr = "localhost:4532"

var ErrUnexpectedValue = fmt.Errorf("Unexpected value in response")

// Rig represents a receiver or tranceiver.
//
// It holds the tcp connection to the service (rigctld).
type TCPRig struct {
	mu   sync.Mutex
	conn *textproto.Conn
	addr string
}

// VFO (Variable Frequency Oscillator) represents a tunable channel,
// from the radio operator's view.
//
// Also referred to as "BAND" (A-band/B-band) by some radio manufacturers.
type tcpVFO struct{ r *TCPRig }

// OpenTCP connects to the rigctld service and returns a ready to use Rig.
//
// Caller must remember to Close the Rig after use.
func OpenTCP(addr string) (*TCPRig, error) {
	r := &TCPRig{addr: addr}
	return r, r.dial()
}

func (r *TCPRig) dial() (err error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.conn != nil {
		r.conn.Close()
	}

	r.conn, err = textproto.Dial("tcp", r.addr)

	return err
}

// Closes the connection to the Rig.
func (r *TCPRig) Close() error {
	if r.conn == nil {
		return nil
	}
	return r.conn.Close()
}

// Returns the Rig's active VFO (for control).
func (r *TCPRig) CurrentVFO() VFO { return &tcpVFO{r} }

// Gets the dial frequency for this VFO.
func (v *tcpVFO) GetFreq() (int, error) {
	resp, err := v.cmd("f")
	if err != nil {
		return -1, err
	}

	freq, err := strconv.Atoi(resp)
	if err != nil {
		return -1, err
	}

	return freq, nil
}

// Sets the dial frequency for this VFO.
func (v *tcpVFO) SetFreq(freq int) error {
	_, err := v.cmd("F %d", freq)
	return err
}

// GetPTT returns the PTT state for this VFO.
func (v *tcpVFO) GetPTT() (bool, error) {
	resp, err := v.cmd("t")
	if err != nil {
		return false, err
	}

	switch resp {
	case "0":
		return false, nil
	case "1":
		return true, nil
	default:
		return false, ErrUnexpectedValue
	}
}

// Enable (or disable) PTT on this VFO.
func (v *tcpVFO) SetPTT(on bool) error {
	bInt := 0
	if on == true {
		bInt = 1
	}

	_, err := v.cmd("t %d", bInt)
	return err
}

// TODO: Move retry logic to *TCPRig
func (v *tcpVFO) cmd(format string, args ...interface{}) (string, error) {
	var err error
	var resp string

	// Retry
	for i := 0; i < 3; i++ {
		if v.r.conn == nil {
			// Try re-dialing
			if err = v.r.dial(); err != nil {
				break
			}
		}

		resp, err = v.r.cmd(format, args...)
		if err == nil {
			break
		}

		_, isNetError := err.(net.Error)
		if err == io.EOF || isNetError {
			v.r.conn = nil
		}

	}

	return resp, err
}

func (r *TCPRig) cmd(format string, args ...interface{}) (string, error) {
	id, err := r.conn.Cmd(format, args...)
	if err != nil {
		return "", err
	}

	r.conn.StartResponse(id)
	defer r.conn.EndResponse(id)

	resp, err := r.conn.ReadLine()
	if err != nil {
		return "", err
	} else if err := toError(resp); err != nil {
		return resp, err
	}

	return resp, nil
}

func toError(str string) error {
	if !strings.HasPrefix(str, "RPRT ") {
		return nil
	}

	parts := strings.SplitN(str, " ", 2)

	code, err := strconv.Atoi(parts[1])
	if err != nil {
		return err
	}

	switch code {
	case 0:
		return nil
	default:
		return fmt.Errorf("code %d", code)
	}
}
<commit_msg>hamlib: Add TCP io timeouts<commit_after>// Copyright 2015 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

package hamlib

import (
	"fmt"
	"io"
	"net"
	"net/textproto"
	"strconv"
	"strings"
	"sync"
	"time"
)

const DefaultTCPAddr = "localhost:4532"

var ErrUnexpectedValue = fmt.Errorf("Unexpected value in response")

// TCPTimeout defines the timeout duration of dial, read and write operations.
var TCPTimeout = time.Second

// Rig represents a receiver or tranceiver.
//
// It holds the tcp connection to the service (rigctld).
type TCPRig struct {
	mu      sync.Mutex
	conn    *textproto.Conn
	tcpConn net.Conn
	addr    string
}

// VFO (Variable Frequency Oscillator) represents a tunable channel,
// from the radio operator's view.
//
// Also referred to as "BAND" (A-band/B-band) by some radio manufacturers.
type tcpVFO struct{ r *TCPRig }

// OpenTCP connects to the rigctld service and returns a ready to use Rig.
//
// Caller must remember to Close the Rig after use.
func OpenTCP(addr string) (*TCPRig, error) {
	r := &TCPRig{addr: addr}
	return r, r.dial()
}

func (r *TCPRig) dial() (err error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.conn != nil {
		r.conn.Close()
	}

	// Dial with 3 second timeout
	r.tcpConn, err = net.DialTimeout("tcp", r.addr, TCPTimeout)
	if err != nil {
		return err
	}

	r.conn = textproto.NewConn(r.tcpConn)

	return err
}

// Closes the connection to the Rig.
func (r *TCPRig) Close() error {
	if r.conn == nil {
		return nil
	}
	return r.conn.Close()
}

// Returns the Rig's active VFO (for control).
func (r *TCPRig) CurrentVFO() VFO { return &tcpVFO{r} }

// Gets the dial frequency for this VFO.
func (v *tcpVFO) GetFreq() (int, error) {
	resp, err := v.cmd("f")
	if err != nil {
		return -1, err
	}

	freq, err := strconv.Atoi(resp)
	if err != nil {
		return -1, err
	}

	return freq, nil
}

// Sets the dial frequency for this VFO.
func (v *tcpVFO) SetFreq(freq int) error {
	_, err := v.cmd("F %d", freq)
	return err
}

// GetPTT returns the PTT state for this VFO.
func (v *tcpVFO) GetPTT() (bool, error) {
	resp, err := v.cmd("t")
	if err != nil {
		return false, err
	}

	switch resp {
	case "0":
		return false, nil
	case "1":
		return true, nil
	default:
		return false, ErrUnexpectedValue
	}
}

// Enable (or disable) PTT on this VFO.
func (v *tcpVFO) SetPTT(on bool) error {
	bInt := 0
	if on == true {
		bInt = 1
	}

	_, err := v.cmd("t %d", bInt)
	return err
}

// TODO: Move retry logic to *TCPRig
func (v *tcpVFO) cmd(format string, args ...interface{}) (string, error) {
	var err error
	var resp string

	// Retry
	for i := 0; i < 3; i++ {
		if v.r.conn == nil {
			// Try re-dialing
			if err = v.r.dial(); err != nil {
				break
			}
		}

		resp, err = v.r.cmd(format, args...)
		if err == nil {
			break
		}

		_, isNetError := err.(net.Error)
		if err == io.EOF || isNetError {
			v.r.conn = nil
		}

	}

	return resp, err
}

func (r *TCPRig) cmd(format string, args ...interface{}) (string, error) {
	r.tcpConn.SetDeadline(time.Now().Add(TCPTimeout))
	id, err := r.conn.Cmd(format, args...)
	r.tcpConn.SetDeadline(time.Time{})

	if err != nil {
		return "", err
	}

	r.conn.StartResponse(id)
	defer r.conn.EndResponse(id)

	r.tcpConn.SetDeadline(time.Now().Add(TCPTimeout))
	resp, err := r.conn.ReadLine()
	r.tcpConn.SetDeadline(time.Time{})

	if err != nil {
		return "", err
	} else if err := toError(resp); err != nil {
		return resp, err
	}

	return resp, nil
}

func toError(str string) error {
	if !strings.HasPrefix(str, "RPRT ") {
		return nil
	}

	parts := strings.SplitN(str, " ", 2)

	code, err := strconv.Atoi(parts[1])
	if err != nil {
		return err
	}

	switch code {
	case 0:
		return nil
	default:
		return fmt.Errorf("code %d", code)
	}
}
<|endoftext|>


<commit_before>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

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

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

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

package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/percona/cloud-protocol/proto"
	"io/ioutil"
	golog "log"
	"net/http"
	"os"
	"strings"
	"time"
)

const (
	VERSION = "1.0.0"
)

func init() {
	golog.SetFlags(golog.Ldate | golog.Ltime | golog.Lmicroseconds | golog.Lshortfile)
}

type Cli struct {
	// State
	apiHostname string
	apiKey      string
	connected   bool
	agentUuid   string
	client      *http.Client
	entryLinks  map[string]string
	agentLinks  map[string]string
}

func main() {
	cli := &Cli{}
	cli.Run()
}

func (cli *Cli) Run() {
	fmt.Printf("percona-agent-cli %s\nType '?' for help.\nUse 'connect' to get started.\n\n", VERSION)

	bio := bufio.NewReader(os.Stdin)
	for {
		fmt.Printf("%s@%s> ", cli.agentUuid, cli.apiHostname)
		line, _, err := bio.ReadLine()
		if err != nil {
			golog.Println(err)
			continue
		}
		lines := strings.Split(string(line), ";")
		for _, line := range lines {
			args := strings.Split(strings.TrimSpace(line), " ")
			if len(args) == 0 {
				continue
			}
			cli.doCmd(args)
		}
	}
}

func (cli *Cli) doCmd(args []string) {
	t0 := time.Now()

	switch args[0] {
	case "":
		return
	case "?", "help":
		cli.help()
		return
	case "connect":
		cli.connect(args)
	case "agent":
		cli.agent(args)
	case "status":
		cli.status(args)
	case "config":
		cli.config(args)
	case "send":
		cli.send(args)
	case "info":
		cli.info(args)
	default:
		fmt.Println("Unknown command: " + args[0])
		return
	}

	d := time.Now().Sub(t0)
	fmt.Printf("%s %s\n", args[0], d)
}

func (cli *Cli) help() {
	fmt.Printf("Commands:\n  connect\n  agent\n  status\n  ?\n\n")
	fmt.Printf("Prompt:\n  agent@api>\n  Use 'connect' command to connect to API, then 'agent' command to set agent.\n\n")
	fmt.Printf("CTRL-C to exit\n\n")
}

func (cli *Cli) connect(args []string) {
	if cli.client == nil {
		cli.client = &http.Client{}
	}

	if len(args) != 3 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 3\n", len(args))
		fmt.Println("Usage: connect api-hostname api-key")
		fmt.Println("Exmaple: connect http://localhost:8000 00000000000000000000000000000001")
		return
	}

	apiHostname := args[1]
	cli.apiKey = args[2] // set now because Get() uses it

	data := cli.Get(apiHostname)
	links := &proto.Links{}
	if err := json.Unmarshal(data, links); err != nil {
		golog.Printf("GET %s error: json.Unmarshal: %s: %s", apiHostname, err, string(data))
		return
	}

	if _, ok := links.Links["agents"]; !ok {
		fmt.Println("ERROR: Connected but no agents link.  Try to connect again.  API has bug if problem continues.")
		return
	}

	cli.apiHostname = apiHostname
	cli.entryLinks = links.Links
	cli.connected = true

	fmt.Printf("Entry links:\n%+v\n\n", cli.entryLinks)
}

func (cli *Cli) agent(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}

	if len(args) == 1 && cli.agentUuid != "" {
		cli.agentUuid = ""
		cli.agentLinks = make(map[string]string)
		return
	}

	if len(args) != 2 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 2\n", len(args))
		fmt.Println("Usage: agent agent-uuid")
		fmt.Println("Exmaple: agent 00000000-0000-0000-0000-000000000001")
		return
	}

	uuid := args[1]

	url := cli.entryLinks["agents"] + "/" + uuid
	data := cli.Get(url)
	links := &proto.Links{}
	if err := json.Unmarshal(data, links); err != nil {
		golog.Printf("GET %s error: json.Unmarshal: %s: %s", url, err, string(data))
		return
	}

	needLinks := []string{"self", "cmd", "log", "data"}
	for _, needLink := range needLinks {
		if _, ok := links.Links[needLink]; !ok {
			fmt.Println("ERROR: API did not return a %s link.  Reconnect and try again.\n", needLink)
			return
		}
	}

	cli.agentUuid = uuid
	cli.agentLinks = links.Links
	fmt.Printf("Agent links:\n%+v\n\n", cli.agentLinks)
}

func (cli *Cli) status(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}
	if cli.agentUuid == "" {
		fmt.Println("Agent UUID not set.  Use 'agent' command.")
		return
	}
	status := cli.Get(cli.agentLinks["self"] + "/status")
	fmt.Println(string(status))
}

func (cli *Cli) config(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}
	if cli.agentUuid == "" {
		fmt.Println("Agent UUID not set.  Use 'agent' command.")
		return
	}

	if len(args) != 3 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 3\n", len(args))
		fmt.Println("Usage: config update file")
		fmt.Println("Exmaple: config update /tmp/new-log.conf")
		return
	}

	if args[1] != "update" {
		fmt.Printf("Invalid arg: got %s, expected 'config'n", args[1])
		return
	}

	// todo
	_, err := ioutil.ReadFile(args[2])
	if err != nil {
		golog.Println(err)
		return
	}
}

func (cli *Cli) send(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}
	if cli.agentUuid == "" {
		fmt.Println("Agent UUID not set.  Use 'agent' command.")
		return
	}
	if len(args) < 3 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 3\n", len(args))
		fmt.Println("Usage: send cmd service")
		fmt.Println("Exmaple: send Stop agent")
		return
	}
	cmd := &proto.Cmd{
		Ts:        time.Now(),
		User:      "percona-agent-cli",
		AgentUuid: cli.agentUuid,
		Cmd:       args[1],
		Service:   args[2],
	}
	if len(args) == 4 {
		switch args[1] {
		case "Update":
			cmd.Data = []byte(args[3])
		case "GetInfo":
			si := &proto.ServiceInstance{}
			if err := json.Unmarshal([]byte(args[3]), si); err != nil {
				fmt.Printf("ERROR: %s\n", err)
				return
			}
			bytes, _ := json.Marshal(si)
			cmd.Data = bytes
		}
	}
	fmt.Printf("%#v\n", cmd)
	reply, err := cli.Put(cli.agentLinks["self"]+"/cmd", cmd)
	if err != nil {
		golog.Println(err)
		return
	}
	if reply.Error != "" {
		fmt.Printf("ERROR: %s\n", reply.Error)
		return
	}
	fmt.Println("OK")
	switch cmd.Cmd {
	case "Version":
		v := &proto.Version{}
		if err := json.Unmarshal(reply.Data, v); err != nil {
			fmt.Printf("Invalid Version reply: %s\n", err)
			return
		}
		fmt.Printf("%#v\n", v)
	}
}

func (cli *Cli) info(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}
	if cli.agentUuid == "" {
		fmt.Println("Agent UUID not set.  Use 'agent' command.")
		return
	}
	if len(args) < 2 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 2\n", len(args))
		fmt.Println("Usage: info service [data]")
		fmt.Println("Exmaple: info mysql user:pass@tcp")
		return
	}
	cmd := &proto.Cmd{
		Ts:        time.Now(),
		User:      "percona-agent-cli",
		AgentUuid: cli.agentUuid,
		Cmd:       "GetInfo",
		Service:   "instance",
	}
	if len(args) == 3 {
		switch args[1] {
		case "mysql":
			mi := &proto.MySQLInstance{
				DSN: args[2],
			}
			bytes, err := json.Marshal(mi)
			if err != nil {
				fmt.Printf("ERROR: %s\n", err)
				return
			}
			si := &proto.ServiceInstance{
				Service: "mysql",
				Instance: bytes,
			}
			bytes, err = json.Marshal(si)
			if err != nil {
				fmt.Printf("ERROR: %s\n", err)
				return
			}
			cmd.Data = bytes
		}
	}
	reply, err := cli.Put(cli.agentLinks["self"]+"/cmd", cmd)
	if err != nil {
		golog.Println(err)
		return
	}
	if reply.Error != "" {
		fmt.Printf("ERROR: %s\n", reply.Error)
		return
	}
	fmt.Println("OK")
	switch args[1] {
	case "mysql":
		mi := &proto.MySQLInstance{}
		if err := json.Unmarshal(reply.Data, mi); err != nil {
			fmt.Printf("Invalid reply: %s\n", err)
			return
		}
		fmt.Printf("%#v\n", mi)
	}
}

func (cli *Cli) Get(url string) []byte {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		golog.Fatal(err)
	}
	req.Header.Add("X-Percona-API-Key", cli.apiKey)

	resp, err := cli.client.Do(req)
	if err != nil {
		golog.Printf("GET %s error: client.Do: %s", url, err)
		return nil
	}
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		golog.Printf("GET %s error: ioutil.ReadAll: %s", url, err)
		return nil
	}
	return body
}

func (cli *Cli) Put(url string, cmd *proto.Cmd) (*proto.Reply, error) {
	golog.Printf("POST %s\n", url)
	data, err := json.Marshal(cmd)
	if err != nil {
		return nil, err
	}
	buf := bytes.NewBuffer(data)
	req, err := http.NewRequest("PUT", url, buf)
	if err != nil {
		return nil, err
	}
	req.Header.Add("X-Percona-API-Key", cli.apiKey)
	resp, err := cli.client.Do(req)
	if err != nil {
		return nil, err
	}
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return nil, err
	}
	reply := &proto.Reply{}
	if err := json.Unmarshal(body, reply); err != nil {
		return nil, err
	}
	return reply, nil
}
<commit_msg>Add GetInfo server support to cli.<commit_after>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

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

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

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

package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/percona/cloud-protocol/proto"
	"io/ioutil"
	golog "log"
	"net/http"
	"os"
	"strings"
	"time"
)

func init() {
	golog.SetFlags(golog.Ldate | golog.Ltime | golog.Lmicroseconds | golog.Lshortfile)
}

type Cli struct {
	// State
	apiHostname string
	apiKey      string
	connected   bool
	agentUuid   string
	client      *http.Client
	entryLinks  map[string]string
	agentLinks  map[string]string
}

func main() {
	cli := &Cli{}
	cli.Run()
}

func (cli *Cli) Run() {
	fmt.Printf("percona-agent-cli\nType '?' for help.\nUse 'connect' to get started.\n\n")

	bio := bufio.NewReader(os.Stdin)
	for {
		fmt.Printf("%s@%s> ", cli.agentUuid, cli.apiHostname)
		line, _, err := bio.ReadLine()
		if err != nil {
			golog.Println(err)
			continue
		}
		lines := strings.Split(string(line), ";")
		for _, line := range lines {
			args := strings.Split(strings.TrimSpace(line), " ")
			if len(args) == 0 {
				continue
			}
			cli.doCmd(args)
		}
	}
}

func (cli *Cli) doCmd(args []string) {
	t0 := time.Now()

	switch args[0] {
	case "":
		return
	case "?", "help":
		cli.help()
		return
	case "connect":
		cli.connect(args)
	case "agent":
		cli.agent(args)
	case "status":
		cli.status(args)
	case "config":
		cli.config(args)
	case "send":
		cli.send(args)
	case "info":
		cli.info(args)
	default:
		fmt.Println("Unknown command: " + args[0])
		return
	}

	d := time.Now().Sub(t0)
	fmt.Printf("%s %s\n", args[0], d)
}

func (cli *Cli) help() {
	fmt.Printf("Commands:\n  connect\n  agent\n  status\n  ?\n\n")
	fmt.Printf("Prompt:\n  agent@api>\n  Use 'connect' command to connect to API, then 'agent' command to set agent.\n\n")
	fmt.Printf("CTRL-C to exit\n\n")
}

func (cli *Cli) connect(args []string) {
	if cli.client == nil {
		cli.client = &http.Client{}
	}

	if len(args) != 3 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 3\n", len(args))
		fmt.Println("Usage: connect api-hostname api-key")
		fmt.Println("Exmaple: connect http://localhost:8000 00000000000000000000000000000001")
		return
	}

	apiHostname := args[1]
	cli.apiKey = args[2] // set now because Get() uses it

	data := cli.Get(apiHostname)
	links := &proto.Links{}
	if err := json.Unmarshal(data, links); err != nil {
		golog.Printf("GET %s error: json.Unmarshal: %s: %s", apiHostname, err, string(data))
		return
	}

	if _, ok := links.Links["agents"]; !ok {
		fmt.Println("ERROR: Connected but no agents link.  Try to connect again.  API has bug if problem continues.")
		return
	}

	cli.apiHostname = apiHostname
	cli.entryLinks = links.Links
	cli.connected = true

	fmt.Printf("Entry links:\n%+v\n\n", cli.entryLinks)
}

func (cli *Cli) agent(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}

	if len(args) == 1 && cli.agentUuid != "" {
		cli.agentUuid = ""
		cli.agentLinks = make(map[string]string)
		return
	}

	if len(args) != 2 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 2\n", len(args))
		fmt.Println("Usage: agent agent-uuid")
		fmt.Println("Exmaple: agent 00000000-0000-0000-0000-000000000001")
		return
	}

	uuid := args[1]

	url := cli.entryLinks["agents"] + "/" + uuid
	data := cli.Get(url)
	links := &proto.Links{}
	if err := json.Unmarshal(data, links); err != nil {
		golog.Printf("GET %s error: json.Unmarshal: %s: %s", url, err, string(data))
		return
	}

	needLinks := []string{"self", "cmd", "log", "data"}
	for _, needLink := range needLinks {
		if _, ok := links.Links[needLink]; !ok {
			fmt.Println("ERROR: API did not return a %s link.  Reconnect and try again.\n", needLink)
			return
		}
	}

	cli.agentUuid = uuid
	cli.agentLinks = links.Links
	fmt.Printf("Agent links:\n%+v\n\n", cli.agentLinks)
}

func (cli *Cli) status(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}
	if cli.agentUuid == "" {
		fmt.Println("Agent UUID not set.  Use 'agent' command.")
		return
	}
	status := cli.Get(cli.agentLinks["self"] + "/status")
	fmt.Println(string(status))
}

func (cli *Cli) config(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}
	if cli.agentUuid == "" {
		fmt.Println("Agent UUID not set.  Use 'agent' command.")
		return
	}

	if len(args) != 3 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 3\n", len(args))
		fmt.Println("Usage: config update file")
		fmt.Println("Exmaple: config update /tmp/new-log.conf")
		return
	}

	if args[1] != "update" {
		fmt.Printf("Invalid arg: got %s, expected 'config'n", args[1])
		return
	}

	// todo
	_, err := ioutil.ReadFile(args[2])
	if err != nil {
		golog.Println(err)
		return
	}
}

func (cli *Cli) send(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}
	if cli.agentUuid == "" {
		fmt.Println("Agent UUID not set.  Use 'agent' command.")
		return
	}
	if len(args) < 3 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 3\n", len(args))
		fmt.Println("Usage: send cmd service")
		fmt.Println("Exmaple: send Stop agent")
		return
	}
	cmd := &proto.Cmd{
		Ts:        time.Now(),
		User:      "percona-agent-cli",
		AgentUuid: cli.agentUuid,
		Cmd:       args[1],
		Service:   args[2],
	}
	if len(args) == 4 {
		switch args[1] {
		case "Update":
			cmd.Data = []byte(args[3])
		default:
			fmt.Printf("Unknown arg: %s\n", args[3])
			return
		}
	}
	fmt.Printf("%#v\n", cmd)
	reply, err := cli.Put(cli.agentLinks["self"]+"/cmd", cmd)
	if err != nil {
		golog.Println(err)
		return
	}
	if reply.Error != "" {
		fmt.Printf("ERROR: %s\n", reply.Error)
		return
	}
	fmt.Println("OK")
	switch cmd.Cmd {
	case "Version":
		v := &proto.Version{}
		if err := json.Unmarshal(reply.Data, v); err != nil {
			fmt.Printf("Invalid Version reply: %s\n", err)
			return
		}
		fmt.Printf("%#v\n", v)
	}
}

func (cli *Cli) info(args []string) {
	if !cli.connected {
		fmt.Println("Not connected to API.  Use 'connect' command.")
		return
	}
	if cli.agentUuid == "" {
		fmt.Println("Agent UUID not set.  Use 'agent' command.")
		return
	}
	if len(args) < 2 {
		fmt.Printf("ERROR: Invalid number of args: got %d, expected 2\n", len(args))
		fmt.Println("Usage: info service [data]")
		fmt.Println("Exmaple: info mysql user:pass@tcp")
		return
	}
	cmd := &proto.Cmd{
		Ts:        time.Now(),
		User:      "percona-agent-cli",
		AgentUuid: cli.agentUuid,
		Cmd:       "GetInfo",
		Service:   "instance",
	}
	if len(args) == 3 {
		switch args[1] {
		case "mysql":
			mi := &proto.MySQLInstance{
				DSN: args[2],
			}
			bytes, err := json.Marshal(mi)
			if err != nil {
				fmt.Printf("ERROR: %s\n", err)
				return
			}
			si := &proto.ServiceInstance{
				Service: "mysql",
				Instance: bytes,
			}
			bytes, err = json.Marshal(si)
			if err != nil {
				fmt.Printf("ERROR: %s\n", err)
				return
			}
			cmd.Data = bytes
		case "server":
			mi := &proto.ServerInstance{
				Hostname: args[2],
			}
			bytes, err := json.Marshal(mi)
			if err != nil {
				fmt.Printf("ERROR: %s\n", err)
				return
			}
			si := &proto.ServiceInstance{
				Service: "server",
				Instance: bytes,
			}
			bytes, err = json.Marshal(si)
			if err != nil {
				fmt.Printf("ERROR: %s\n", err)
				return
			}
			cmd.Data = bytes
		}
	}
	reply, err := cli.Put(cli.agentLinks["self"]+"/cmd", cmd)
	if err != nil {
		golog.Println(err)
		return
	}
	if reply.Error != "" {
		fmt.Printf("ERROR: %s\n", reply.Error)
		return
	}
	fmt.Println("OK")
	switch args[1] {
	case "mysql":
		mi := &proto.MySQLInstance{}
		if err := json.Unmarshal(reply.Data, mi); err != nil {
			fmt.Printf("Invalid reply: %s\n", err)
			return
		}
		fmt.Printf("%#v\n", mi)
	case "server":
		mi := &proto.ServerInstance{}
		if err := json.Unmarshal(reply.Data, mi); err != nil {
			fmt.Printf("Invalid reply: %s\n", err)
			return
		}
		fmt.Printf("%#v\n", mi)
	}
}

func (cli *Cli) Get(url string) []byte {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		golog.Fatal(err)
	}
	req.Header.Add("X-Percona-API-Key", cli.apiKey)

	resp, err := cli.client.Do(req)
	if err != nil {
		golog.Printf("GET %s error: client.Do: %s", url, err)
		return nil
	}
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		golog.Printf("GET %s error: ioutil.ReadAll: %s", url, err)
		return nil
	}
	return body
}

func (cli *Cli) Put(url string, cmd *proto.Cmd) (*proto.Reply, error) {
	golog.Printf("POST %s\n", url)
	data, err := json.Marshal(cmd)
	if err != nil {
		return nil, err
	}
	buf := bytes.NewBuffer(data)
	req, err := http.NewRequest("PUT", url, buf)
	if err != nil {
		return nil, err
	}
	req.Header.Add("X-Percona-API-Key", cli.apiKey)
	resp, err := cli.client.Do(req)
	if err != nil {
		return nil, err
	}
	body, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return nil, err
	}
	reply := &proto.Reply{}
	if err := json.Unmarshal(body, reply); err != nil {
		return nil, err
	}
	return reply, nil
}
<|endoftext|>


<commit_before>package account_test

import (
    "os"
    "testing"
    "github.com/pdepip/go-binance/binance"
)

func TestGetTrades(t *testing.T) {
    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    trades, err := client.GetTrades("BNBETH")

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

    t.Logf("%+v\n", trades)
}

func TestGetWithdraws(t *testing.T) {
    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    withdraws, err := client.GetWithdrawHistory()

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

    t.Logf("%+v\n", withdraws)
}

func TestGetDeposits(t *testing.T) {
    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    deposits, err := client.GetDepositHistory()

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

    t.Logf("%+v\n", deposits)
}

func TestGetTradesFromOrder(t *testing.T) {
    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    trades, err := client.GetTradesFromOrder("LINKETH", 10107102)
    if err != nil {
        t.Fatal(err)
    }

    t.Logf("%+v\n", trades)
}

/*

func TestGetPositions(t *testing.T) {

    binance := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    positions, err := binance.GetPositions()

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

    t.Logf("%+v\n", positions)
}

    THE FOLLOWING TESTS WILL PLACE ACTUAL ORDERS ON THE BINANCE EXCHANGE

func TestLimitOrder(t *testing.T) {

    // Params
    order := binance.LimitOrder {
        Symbol:      "BNBBTC",
        Side:        "BUY",
        Type:        "LIMIT",
        TimeInForce: "GTC",
        Quantity:    50.0,
        Price:       0.00025,
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    res, err := client.PlaceLimitOrder(order)

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

    t.Logf("%+v\n", res.OrderId)

}

func TestMarketOrder(t *testing.T) {

    // Params
    order := binance.MarketOrder {
        Symbol:      "BNBBTC",
        Side:        "BUY",
        Type:        "MARKET",
        Quantity:    10.0,
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    res, err := client.PlaceMarketOrder(order)

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

    t.Logf("%+v\n", res.OrderId)

}

func TestQueryOrder(t *testing.T) {

    // Order Param
    order := binance.LimitOrder {
        Symbol:      "BNBBTC",
        Side:        "BUY",
        Type:        "LIMIT",
        TimeInForce: "GTC",
        Quantity:    50.0,
        Price:       0.00025,
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    res, err := client.PlaceLimitOrder(order)
    if err != nil {
        t.Fatal(err)
    }

    // Query Param
    query := binance.OrderQuery{
        Symbol:  res.Symbol,
        OrderId: res.OrderId,
    }

    status, err := client.CheckOrder(query)

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

    t.Logf("%+v\n", status)
}

func TestCancelOrder(t *testing.T) {

    // Order Param
    order := binance.LimitOrder {
        Symbol:      "BNBBTC",
        Side:        "BUY",
        Type:        "LIMIT",
        TimeInForce: "GTC",
        Quantity:    10.0,
        Price:       0.00025,
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    res, err := client.PlaceLimitOrder(order)
    if err != nil {
        t.Fatal(err)
    }

    // Cancel Param
    query := binance.OrderQuery {
        Symbol: res.Symbol,
        OrderId: res.OrderId,
    }

    canceledOrder, err := client.CancelOrder(query)

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

    t.Logf("%+v\n", canceledOrder)
}

func TestGetOpenOrders(t *testing.T) {

    // Param
    query := binance.OpenOrdersQuery {
        Symbol: "BNBBTC",
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    openOrders, err := client.GetOpenOrders(query)
    if err != nil {
        t.Fatal(err)
    }

    t.Logf("%+v\n", openOrders)
}

*/
<commit_msg>adds test for getaccountinfo<commit_after>package account_test

import (
    "os"
    "testing"
    "github.com/pdepip/go-binance/binance"
)

func TestGetAccountInfo(t *testing.T) {

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    res, err := client.GetAccountInfo()
    if err != nil {
        t.Fatal(err)
    }

    t.Logf("%+v\n", res)
}

func TestGetTrades(t *testing.T) {
    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    trades, err := client.GetTrades("BNBETH")

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

    t.Logf("%+v\n", trades)
}

func TestGetWithdraws(t *testing.T) {
    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    withdraws, err := client.GetWithdrawHistory()

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

    t.Logf("%+v\n", withdraws)
}

func TestGetDeposits(t *testing.T) {
    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    deposits, err := client.GetDepositHistory()

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

    t.Logf("%+v\n", deposits)
}

func TestGetTradesFromOrder(t *testing.T) {
    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))

    trades, err := client.GetTradesFromOrder("LINKETH", 10107102)
    if err != nil {
        t.Fatal(err)
    }

    t.Logf("%+v\n", trades)
}

/*

func TestGetPositions(t *testing.T) {

    binance := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    positions, err := binance.GetPositions()

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

    t.Logf("%+v\n", positions)
}

    THE FOLLOWING TESTS WILL PLACE ACTUAL ORDERS ON THE BINANCE EXCHANGE

func TestLimitOrder(t *testing.T) {

    // Params
    order := binance.LimitOrder {
        Symbol:      "BNBBTC",
        Side:        "BUY",
        Type:        "LIMIT",
        TimeInForce: "GTC",
        Quantity:    50.0,
        Price:       0.00025,
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    res, err := client.PlaceLimitOrder(order)

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

    t.Logf("%+v\n", res.OrderId)

}

func TestMarketOrder(t *testing.T) {

    // Params
    order := binance.MarketOrder {
        Symbol:      "BNBBTC",
        Side:        "BUY",
        Type:        "MARKET",
        Quantity:    10.0,
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    res, err := client.PlaceMarketOrder(order)

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

    t.Logf("%+v\n", res.OrderId)

}

func TestQueryOrder(t *testing.T) {

    // Order Param
    order := binance.LimitOrder {
        Symbol:      "BNBBTC",
        Side:        "BUY",
        Type:        "LIMIT",
        TimeInForce: "GTC",
        Quantity:    50.0,
        Price:       0.00025,
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    res, err := client.PlaceLimitOrder(order)
    if err != nil {
        t.Fatal(err)
    }

    // Query Param
    query := binance.OrderQuery{
        Symbol:  res.Symbol,
        OrderId: res.OrderId,
    }

    status, err := client.CheckOrder(query)

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

    t.Logf("%+v\n", status)
}

func TestCancelOrder(t *testing.T) {

    // Order Param
    order := binance.LimitOrder {
        Symbol:      "BNBBTC",
        Side:        "BUY",
        Type:        "LIMIT",
        TimeInForce: "GTC",
        Quantity:    10.0,
        Price:       0.00025,
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    res, err := client.PlaceLimitOrder(order)
    if err != nil {
        t.Fatal(err)
    }

    // Cancel Param
    query := binance.OrderQuery {
        Symbol: res.Symbol,
        OrderId: res.OrderId,
    }

    canceledOrder, err := client.CancelOrder(query)

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

    t.Logf("%+v\n", canceledOrder)
}

func TestGetOpenOrders(t *testing.T) {

    // Param
    query := binance.OpenOrdersQuery {
        Symbol: "BNBBTC",
    }

    client := binance.New(os.Getenv("BINANCE_KEY"), os.Getenv("BINANCE_SECRET"))
    openOrders, err := client.GetOpenOrders(query)
    if err != nil {
        t.Fatal(err)
    }

    t.Logf("%+v\n", openOrders)
}

*/
<|endoftext|>


<commit_before>package collector

import (
	"fullerite/metric"
	"test_utils"

	"encoding/json"
	"fmt"
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestDiamondConfigureEmptyConfig(t *testing.T) {
	config := make(map[string]interface{})

	d := NewDiamond(nil, 12, nil)
	d.Configure(config)

	assert.Equal(t,
		d.Interval(),
		12,
		"should be the default collection interval",
	)
}

func TestDiamondConfigure(t *testing.T) {
	config := make(map[string]interface{})
	config["interval"] = 9999
	config["port"] = "0"
	d := NewDiamond(nil, 12, nil)
	d.Configure(config)

	assert := assert.New(t)
	assert.Equal(d.Interval(), 9999, "should be the defined interval")
	assert.Equal(d.Port(), "0", "should be the defined port")
}

func TestDiamondCollect(t *testing.T) {
	config := make(map[string]interface{})
	config["port"] = "0"

	testChannel := make(chan metric.Metric)
	testLog := test_utils.BuildLogger()

	d := NewDiamond(testChannel, 123, testLog)
	d.Configure(config)

	// start collecting Diamond metrics
	go d.Collect()

	conn, err := connectToDiamondCollector(d)
	require.Nil(t, err, "should connect")
	require.NotNil(t, conn, "should connect")

	emitTestMetric(conn)

	select {
	case m := <-d.Channel():
		assert.Equal(t, m.Name, "test")
	case <-time.After(1 * time.Second):
		t.Fail()
	}
}

func connectToDiamondCollector(d *Diamond) (net.Conn, error) {
	// emit a Diamond metric
	var (
		conn net.Conn
		err  error
	)
	for retry := 0; retry < 3; retry++ {
		if conn, err = net.DialTimeout("tcp", "localhost:"+d.Port(), 2*time.Second); err == nil {
			break
		}
		time.Sleep(1 * time.Second)
	}
	return conn, err
}

func emitTestMetric(conn net.Conn) {
	m := metric.New("test")
	b, _ := json.Marshal(m)
	fmt.Fprintf(conn, string(b)+"\n")
	fmt.Fprintf(conn, string(b)+"\n")
}
<commit_msg>sleep shorter in diamond collector tests<commit_after>package collector

import (
	"fullerite/metric"
	"test_utils"

	"encoding/json"
	"fmt"
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestDiamondConfigureEmptyConfig(t *testing.T) {
	config := make(map[string]interface{})

	d := NewDiamond(nil, 12, nil)
	d.Configure(config)

	assert.Equal(t,
		d.Interval(),
		12,
		"should be the default collection interval",
	)
}

func TestDiamondConfigure(t *testing.T) {
	config := make(map[string]interface{})
	config["interval"] = 9999
	config["port"] = "0"
	d := NewDiamond(nil, 12, nil)
	d.Configure(config)

	assert := assert.New(t)
	assert.Equal(d.Interval(), 9999, "should be the defined interval")
	assert.Equal(d.Port(), "0", "should be the defined port")
}

func TestDiamondCollect(t *testing.T) {
	config := make(map[string]interface{})
	config["port"] = "0"

	testChannel := make(chan metric.Metric)
	testLog := test_utils.BuildLogger()

	d := NewDiamond(testChannel, 123, testLog)
	d.Configure(config)

	// start collecting Diamond metrics
	go d.Collect()

	conn, err := connectToDiamondCollector(d)
	require.Nil(t, err, "should connect")
	require.NotNil(t, conn, "should connect")

	emitTestMetric(conn)

	select {
	case m := <-d.Channel():
		assert.Equal(t, m.Name, "test")
	case <-time.After(1 * time.Second):
		t.Fail()
	}
}

func connectToDiamondCollector(d *Diamond) (net.Conn, error) {
	// emit a Diamond metric
	var (
		conn net.Conn
		err  error
	)
	for retry := 0; retry < 3; retry++ {
		if conn, err = net.DialTimeout("tcp", "localhost:"+d.Port(), 2*time.Second); err == nil {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}
	return conn, err
}

func emitTestMetric(conn net.Conn) {
	m := metric.New("test")
	b, _ := json.Marshal(m)
	fmt.Fprintf(conn, string(b)+"\n")
	fmt.Fprintf(conn, string(b)+"\n")
}
<|endoftext|>


<commit_before>package account

import (
	"context"
	"encoding/json"

	log "github.com/sirupsen/logrus"
	"github.com/tendermint/tmlibs/db"

	"github.com/bytom/blockchain/query"
	"github.com/bytom/blockchain/signers"
	"github.com/bytom/crypto/sha3pool"
	chainjson "github.com/bytom/encoding/json"
	"github.com/bytom/errors"
	"github.com/bytom/protocol"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/legacy"
)

const (
	//UTXOPreFix is account unspent outputs store db by key with this prefix
	UTXOPreFix = "ACU:"
)

var walletkey = []byte("WAL:")

func accountUTXOKey(name string) []byte {
	return []byte(UTXOPreFix + name)
}

//WalletInfo is base valid block info to handle orphan block rollback
type WalletInfo struct {
	Height uint64
	Hash   bc.Hash
}

//Wallet is related to storing account unspent outputs
type Wallet struct {
	DB db.DB
	WalletInfo
}

//NewWallet return a new wallet instance
func NewWallet(db db.DB) *Wallet {
	w := &Wallet{
		DB: db,
	}

	walletInfo, err := w.GetWalletInfo()
	if err != nil {
		log.WithField("warn", err).Warn("get wallet info")
	}
	w.Height = walletInfo.Height
	w.Hash = walletInfo.Hash
	return w
}

//GetWalletHeight return wallet on current height
func (w *Wallet) GetWalletHeight() uint64 {
	return w.Height
}

//GetWalletInfo return stored wallet info and nil,if error,
//return initial wallet info and err
func (w *Wallet) GetWalletInfo() (WalletInfo, error) {
	var info WalletInfo
	var rawWallet []byte

	if rawWallet = w.DB.Get(walletkey); rawWallet == nil {
		return info, nil
	}

	if err := json.Unmarshal(rawWallet, &w); err != nil {
		return info, err
	}

	return info, nil

}

//WalletUpdate process every valid block and reverse every invalid block which need to rollback
func (m *Manager) WalletUpdate(c *protocol.Chain) {
	var err error
	var block *legacy.Block

	storeBatch := m.wallet.DB.NewBatch()

LOOP:

	for !c.InMainChain(m.wallet.Height, m.wallet.Hash) {
		if block, err = c.GetBlockByHash(&m.wallet.Hash); err != nil {
			log.WithField("err", err).Error("get block by hash")
			return
		}

		//Reverse this block
		m.ReverseAccountUTXOs(&storeBatch, block)
		log.WithField("Height", m.wallet.Height).Info("start rollback this block")

		m.wallet.Height = block.Height - 1
		m.wallet.Hash = block.PreviousBlockHash

	}

	//update wallet info and commit batch write
	m.wallet.commitWalletInfo(&storeBatch)

	block, _ = c.GetBlockByHeight(m.wallet.Height + 1)
	//if we already handled the tail of the chain, we wait
	if block == nil {
		<-c.BlockWaiter(m.wallet.Height + 1)
		if block, err = c.GetBlockByHeight(m.wallet.Height + 1); err != nil {
			log.WithField("err", err).Error("wallet get block by height")
			return
		}
	}

	//if false, means that rollback operation is necessary,then goto LOOP
	if block.PreviousBlockHash == m.wallet.Hash {
		//next loop will save
		m.wallet.Height = block.Height
		m.wallet.Hash = block.Hash()
		m.BuildAccountUTXOs(&storeBatch, block)

		//update wallet info and commit batch write
		m.wallet.commitWalletInfo(&storeBatch)
	}

	//goto next loop
	goto LOOP

}

func (w *Wallet) commitWalletInfo(batch *db.Batch) {
	var info WalletInfo

	info.Height = w.Height
	info.Hash = w.Hash

	rawWallet, err := json.Marshal(info)
	if err != nil {
		log.WithField("err", err).Error("save wallet info")
		return
	}
	//update wallet to db
	(*batch).Set(walletkey, rawWallet)
	//commit to db
	(*batch).Write()
}

//UTXO is a structure about account unspent outputs
type UTXO struct {
	OutputID     []byte
	AssetID      []byte
	Amount       uint64
	AccountID    string
	ProgramIndex uint64
	Program      []byte
	SourceID     []byte
	SourcePos    uint64
	RefData      []byte
	Change       bool
}

var emptyJSONObject = json.RawMessage(`{}`)

// A Saver is responsible for saving an annotated account object.
// for indexing and retrieval.
// If the Core is configured not to provide search services,
// SaveAnnotatedAccount can be a no-op.
type Saver interface {
	SaveAnnotatedAccount(context.Context, *query.AnnotatedAccount) error
}

//Annotated init an annotated account object
func Annotated(a *Account) (*query.AnnotatedAccount, error) {
	aa := &query.AnnotatedAccount{
		ID:     a.ID,
		Alias:  a.Alias,
		Quorum: a.Quorum,
		Tags:   &emptyJSONObject,
	}

	tags, err := json.Marshal(a.Tags)
	if err != nil {
		return nil, err
	}
	if len(tags) > 0 {
		rawTags := json.RawMessage(tags)
		aa.Tags = &rawTags
	}

	path := signers.Path(a.Signer, signers.AccountKeySpace)
	var jsonPath []chainjson.HexBytes
	for _, p := range path {
		jsonPath = append(jsonPath, p)
	}
	for _, xpub := range a.XPubs {
		aa.Keys = append(aa.Keys, &query.AccountKey{
			RootXPub:              xpub,
			AccountXPub:           xpub.Derive(path),
			AccountDerivationPath: jsonPath,
		})
	}
	return aa, nil
}

func (m *Manager) indexAnnotatedAccount(ctx context.Context, a *Account) error {
	if m.indexer == nil {
		return nil
	}
	aa, err := Annotated(a)
	if err != nil {
		return err
	}
	return m.indexer.SaveAnnotatedAccount(ctx, aa)
}

type rawOutput struct {
	OutputID bc.Hash
	bc.AssetAmount
	ControlProgram []byte
	txHash         bc.Hash
	outputIndex    uint32
	sourceID       bc.Hash
	sourcePos      uint64
	refData        bc.Hash
}

type accountOutput struct {
	rawOutput
	AccountID string
	keyIndex  uint64
	change    bool
}

//ReverseAccountUTXOs process the invalid blocks when orphan block rollback
func (m *Manager) ReverseAccountUTXOs(batch *db.Batch, b *legacy.Block) {
	var err error

	//unknow how many spent and retire outputs
	reverseOuts := make([]*rawOutput, 0)

	//handle spent UTXOs
	for _, tx := range b.Transactions {
		for _, inpID := range tx.Tx.InputIDs {
			//spend and retire
			sp, err := tx.Spend(inpID)
			if err != nil {
				continue
			}

			resOut, ok := tx.Entries[*sp.SpentOutputId].(*bc.Output)
			if !ok {
				continue
			}

			out := &rawOutput{
				OutputID:       *sp.SpentOutputId,
				AssetAmount:    *resOut.Source.Value,
				ControlProgram: resOut.ControlProgram.Code,
				txHash:         tx.ID,
				sourceID:       *resOut.Source.Ref,
				sourcePos:      resOut.Source.Position,
				refData:        *resOut.Data,
			}
			reverseOuts = append(reverseOuts, out)
		}
	}

	accOuts := m.loadAccountInfo(reverseOuts)
	if err = m.upsertConfirmedAccountOutputs(accOuts, b, batch); err != nil {
		log.WithField("err", err).Error("reversing account spent and retire outputs")
		return
	}

	//handle new UTXOs
	for _, tx := range b.Transactions {
		for j := range tx.Outputs {
			resOutID := tx.ResultIds[j]
			if _, ok := tx.Entries[*resOutID].(*bc.Output); !ok {
				//retirement
				continue
			}
			//delete new UTXOs
			(*batch).Delete(accountUTXOKey(string(resOutID.Bytes())))
		}
	}

}

//BuildAccountUTXOs process valid blocks to build account unspent outputs db
func (m *Manager) BuildAccountUTXOs(batch *db.Batch, b *legacy.Block) {
	var err error

	//handle spent UTXOs
	delOutputIDs := prevoutDBKeys(b.Transactions...)
	for _, delOutputID := range delOutputIDs {
		(*batch).Delete(accountUTXOKey(string(delOutputID.Bytes())))
	}

	//handle new UTXOs
	outs := make([]*rawOutput, 0, len(b.Transactions))
	for _, tx := range b.Transactions {
		for j, out := range tx.Outputs {
			resOutID := tx.ResultIds[j]
			resOut, ok := tx.Entries[*resOutID].(*bc.Output)
			if !ok {
				continue
			}
			out := &rawOutput{
				OutputID:       *tx.OutputID(j),
				AssetAmount:    out.AssetAmount,
				ControlProgram: out.ControlProgram,
				txHash:         tx.ID,
				outputIndex:    uint32(j),
				sourceID:       *resOut.Source.Ref,
				sourcePos:      resOut.Source.Position,
				refData:        *resOut.Data,
			}
			outs = append(outs, out)
		}
	}
	accOuts := m.loadAccountInfo(outs)

	if err = m.upsertConfirmedAccountOutputs(accOuts, b, batch); err != nil {
		log.WithField("err", err).Error("building new account outputs")
		return
	}
}

func prevoutDBKeys(txs ...*legacy.Tx) (outputIDs []bc.Hash) {
	for _, tx := range txs {
		for _, inpID := range tx.Tx.InputIDs {
			if sp, err := tx.Spend(inpID); err == nil {
				outputIDs = append(outputIDs, *sp.SpentOutputId)
			}
		}
	}
	return
}

// loadAccountInfo turns a set of output IDs into a set of
// outputs by adding account annotations.  Outputs that can't be
// annotated are excluded from the result.
func (m *Manager) loadAccountInfo(outs []*rawOutput) []*accountOutput {
	outsByScript := make(map[string][]*rawOutput, len(outs))
	for _, out := range outs {
		scriptStr := string(out.ControlProgram)
		outsByScript[scriptStr] = append(outsByScript[scriptStr], out)
	}

	result := make([]*accountOutput, 0, len(outs))
	cp := controlProgram{}

	var hash []byte
	for s := range outsByScript {
		sha3pool.Sum256(hash, []byte(s))
		bytes := m.db.Get(accountCPKey(string(hash)))
		if bytes == nil {
			continue
		}

		err := json.Unmarshal(bytes, &cp)
		if err != nil {
			continue
		}

		//filte the accounts which exists in accountdb with wallet enabled
		//TODO:filte receiver UTXO about self ?
		isExist := m.db.Get(accountKey(cp.AccountID))
		if isExist == nil {
			continue
		}

		for _, out := range outsByScript[s] {
			newOut := &accountOutput{
				rawOutput: *out,
				AccountID: cp.AccountID,
				keyIndex:  cp.KeyIndex,
				change:    cp.Change,
			}
			result = append(result, newOut)
		}
	}

	return result
}

// upsertConfirmedAccountOutputs records the account data for confirmed utxos.
// If the account utxo already exists (because it's from a local tx), the
// block confirmation data will in the row will be updated.
func (m *Manager) upsertConfirmedAccountOutputs(outs []*accountOutput, block *legacy.Block, batch *db.Batch) error {
	var u *UTXO

	for _, out := range outs {
		u = &UTXO{OutputID: out.OutputID.Bytes(),
			AssetID:      out.AssetId.Bytes(),
			Amount:       out.Amount,
			AccountID:    out.AccountID,
			ProgramIndex: out.keyIndex,
			Program:      out.ControlProgram,
			SourceID:     out.sourceID.Bytes(),
			SourcePos:    out.sourcePos,
			RefData:      out.refData.Bytes(),
			Change:       out.change}

		rawUTXO, err := json.Marshal(u)
		if err != nil {
			return errors.Wrap(err, "failed marshal accountutxo")
		}

		(*batch).Set(accountUTXOKey(string(u.OutputID)), rawUTXO)
	}
	return nil
}
<commit_msg>Update walletInfo key<commit_after>package account

import (
	"context"
	"encoding/json"

	log "github.com/sirupsen/logrus"
	"github.com/tendermint/tmlibs/db"

	"github.com/bytom/blockchain/query"
	"github.com/bytom/blockchain/signers"
	"github.com/bytom/crypto/sha3pool"
	chainjson "github.com/bytom/encoding/json"
	"github.com/bytom/errors"
	"github.com/bytom/protocol"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/legacy"
)

const (
	//UTXOPreFix is account unspent outputs store db by key with this prefix
	UTXOPreFix = "ACU:"
)

var walletkey = []byte("walletInfo")

func accountUTXOKey(name string) []byte {
	return []byte(UTXOPreFix + name)
}

//WalletInfo is base valid block info to handle orphan block rollback
type WalletInfo struct {
	Height uint64
	Hash   bc.Hash
}

//Wallet is related to storing account unspent outputs
type Wallet struct {
	DB db.DB
	WalletInfo
}

//NewWallet return a new wallet instance
func NewWallet(db db.DB) *Wallet {
	w := &Wallet{
		DB: db,
	}

	walletInfo, err := w.GetWalletInfo()
	if err != nil {
		log.WithField("warn", err).Warn("get wallet info")
	}
	w.Height = walletInfo.Height
	w.Hash = walletInfo.Hash
	return w
}

//GetWalletHeight return wallet on current height
func (w *Wallet) GetWalletHeight() uint64 {
	return w.Height
}

//GetWalletInfo return stored wallet info and nil,if error,
//return initial wallet info and err
func (w *Wallet) GetWalletInfo() (WalletInfo, error) {
	var info WalletInfo
	var rawWallet []byte

	if rawWallet = w.DB.Get(walletkey); rawWallet == nil {
		return info, nil
	}

	if err := json.Unmarshal(rawWallet, &w); err != nil {
		return info, err
	}

	return info, nil

}

//WalletUpdate process every valid block and reverse every invalid block which need to rollback
func (m *Manager) WalletUpdate(c *protocol.Chain) {
	var err error
	var block *legacy.Block

	storeBatch := m.wallet.DB.NewBatch()

LOOP:

	for !c.InMainChain(m.wallet.Height, m.wallet.Hash) {
		if block, err = c.GetBlockByHash(&m.wallet.Hash); err != nil {
			log.WithField("err", err).Error("get block by hash")
			return
		}

		//Reverse this block
		m.ReverseAccountUTXOs(&storeBatch, block)
		log.WithField("Height", m.wallet.Height).Info("start rollback this block")

		m.wallet.Height = block.Height - 1
		m.wallet.Hash = block.PreviousBlockHash

	}

	//update wallet info and commit batch write
	m.wallet.commitWalletInfo(&storeBatch)

	block, _ = c.GetBlockByHeight(m.wallet.Height + 1)
	//if we already handled the tail of the chain, we wait
	if block == nil {
		<-c.BlockWaiter(m.wallet.Height + 1)
		if block, err = c.GetBlockByHeight(m.wallet.Height + 1); err != nil {
			log.WithField("err", err).Error("wallet get block by height")
			return
		}
	}

	//if false, means that rollback operation is necessary,then goto LOOP
	if block.PreviousBlockHash == m.wallet.Hash {
		//next loop will save
		m.wallet.Height = block.Height
		m.wallet.Hash = block.Hash()
		m.BuildAccountUTXOs(&storeBatch, block)

		//update wallet info and commit batch write
		m.wallet.commitWalletInfo(&storeBatch)
	}

	//goto next loop
	goto LOOP

}

func (w *Wallet) commitWalletInfo(batch *db.Batch) {
	var info WalletInfo

	info.Height = w.Height
	info.Hash = w.Hash

	rawWallet, err := json.Marshal(info)
	if err != nil {
		log.WithField("err", err).Error("save wallet info")
		return
	}
	//update wallet to db
	(*batch).Set(walletkey, rawWallet)
	//commit to db
	(*batch).Write()
}

//UTXO is a structure about account unspent outputs
type UTXO struct {
	OutputID     []byte
	AssetID      []byte
	Amount       uint64
	AccountID    string
	ProgramIndex uint64
	Program      []byte
	SourceID     []byte
	SourcePos    uint64
	RefData      []byte
	Change       bool
}

var emptyJSONObject = json.RawMessage(`{}`)

// A Saver is responsible for saving an annotated account object.
// for indexing and retrieval.
// If the Core is configured not to provide search services,
// SaveAnnotatedAccount can be a no-op.
type Saver interface {
	SaveAnnotatedAccount(context.Context, *query.AnnotatedAccount) error
}

//Annotated init an annotated account object
func Annotated(a *Account) (*query.AnnotatedAccount, error) {
	aa := &query.AnnotatedAccount{
		ID:     a.ID,
		Alias:  a.Alias,
		Quorum: a.Quorum,
		Tags:   &emptyJSONObject,
	}

	tags, err := json.Marshal(a.Tags)
	if err != nil {
		return nil, err
	}
	if len(tags) > 0 {
		rawTags := json.RawMessage(tags)
		aa.Tags = &rawTags
	}

	path := signers.Path(a.Signer, signers.AccountKeySpace)
	var jsonPath []chainjson.HexBytes
	for _, p := range path {
		jsonPath = append(jsonPath, p)
	}
	for _, xpub := range a.XPubs {
		aa.Keys = append(aa.Keys, &query.AccountKey{
			RootXPub:              xpub,
			AccountXPub:           xpub.Derive(path),
			AccountDerivationPath: jsonPath,
		})
	}
	return aa, nil
}

func (m *Manager) indexAnnotatedAccount(ctx context.Context, a *Account) error {
	if m.indexer == nil {
		return nil
	}
	aa, err := Annotated(a)
	if err != nil {
		return err
	}
	return m.indexer.SaveAnnotatedAccount(ctx, aa)
}

type rawOutput struct {
	OutputID bc.Hash
	bc.AssetAmount
	ControlProgram []byte
	txHash         bc.Hash
	outputIndex    uint32
	sourceID       bc.Hash
	sourcePos      uint64
	refData        bc.Hash
}

type accountOutput struct {
	rawOutput
	AccountID string
	keyIndex  uint64
	change    bool
}

//ReverseAccountUTXOs process the invalid blocks when orphan block rollback
func (m *Manager) ReverseAccountUTXOs(batch *db.Batch, b *legacy.Block) {
	var err error

	//unknow how many spent and retire outputs
	reverseOuts := make([]*rawOutput, 0)

	//handle spent UTXOs
	for _, tx := range b.Transactions {
		for _, inpID := range tx.Tx.InputIDs {
			//spend and retire
			sp, err := tx.Spend(inpID)
			if err != nil {
				continue
			}

			resOut, ok := tx.Entries[*sp.SpentOutputId].(*bc.Output)
			if !ok {
				continue
			}

			out := &rawOutput{
				OutputID:       *sp.SpentOutputId,
				AssetAmount:    *resOut.Source.Value,
				ControlProgram: resOut.ControlProgram.Code,
				txHash:         tx.ID,
				sourceID:       *resOut.Source.Ref,
				sourcePos:      resOut.Source.Position,
				refData:        *resOut.Data,
			}
			reverseOuts = append(reverseOuts, out)
		}
	}

	accOuts := m.loadAccountInfo(reverseOuts)
	if err = m.upsertConfirmedAccountOutputs(accOuts, b, batch); err != nil {
		log.WithField("err", err).Error("reversing account spent and retire outputs")
		return
	}

	//handle new UTXOs
	for _, tx := range b.Transactions {
		for j := range tx.Outputs {
			resOutID := tx.ResultIds[j]
			if _, ok := tx.Entries[*resOutID].(*bc.Output); !ok {
				//retirement
				continue
			}
			//delete new UTXOs
			(*batch).Delete(accountUTXOKey(string(resOutID.Bytes())))
		}
	}

}

//BuildAccountUTXOs process valid blocks to build account unspent outputs db
func (m *Manager) BuildAccountUTXOs(batch *db.Batch, b *legacy.Block) {
	var err error

	//handle spent UTXOs
	delOutputIDs := prevoutDBKeys(b.Transactions...)
	for _, delOutputID := range delOutputIDs {
		(*batch).Delete(accountUTXOKey(string(delOutputID.Bytes())))
	}

	//handle new UTXOs
	outs := make([]*rawOutput, 0, len(b.Transactions))
	for _, tx := range b.Transactions {
		for j, out := range tx.Outputs {
			resOutID := tx.ResultIds[j]
			resOut, ok := tx.Entries[*resOutID].(*bc.Output)
			if !ok {
				continue
			}
			out := &rawOutput{
				OutputID:       *tx.OutputID(j),
				AssetAmount:    out.AssetAmount,
				ControlProgram: out.ControlProgram,
				txHash:         tx.ID,
				outputIndex:    uint32(j),
				sourceID:       *resOut.Source.Ref,
				sourcePos:      resOut.Source.Position,
				refData:        *resOut.Data,
			}
			outs = append(outs, out)
		}
	}
	accOuts := m.loadAccountInfo(outs)

	if err = m.upsertConfirmedAccountOutputs(accOuts, b, batch); err != nil {
		log.WithField("err", err).Error("building new account outputs")
		return
	}
}

func prevoutDBKeys(txs ...*legacy.Tx) (outputIDs []bc.Hash) {
	for _, tx := range txs {
		for _, inpID := range tx.Tx.InputIDs {
			if sp, err := tx.Spend(inpID); err == nil {
				outputIDs = append(outputIDs, *sp.SpentOutputId)
			}
		}
	}
	return
}

// loadAccountInfo turns a set of output IDs into a set of
// outputs by adding account annotations.  Outputs that can't be
// annotated are excluded from the result.
func (m *Manager) loadAccountInfo(outs []*rawOutput) []*accountOutput {
	outsByScript := make(map[string][]*rawOutput, len(outs))
	for _, out := range outs {
		scriptStr := string(out.ControlProgram)
		outsByScript[scriptStr] = append(outsByScript[scriptStr], out)
	}

	result := make([]*accountOutput, 0, len(outs))
	cp := controlProgram{}

	var hash []byte
	for s := range outsByScript {
		sha3pool.Sum256(hash, []byte(s))
		bytes := m.db.Get(accountCPKey(string(hash)))
		if bytes == nil {
			continue
		}

		err := json.Unmarshal(bytes, &cp)
		if err != nil {
			continue
		}

		//filte the accounts which exists in accountdb with wallet enabled
		//TODO:filte receiver UTXO about self ?
		isExist := m.db.Get(accountKey(cp.AccountID))
		if isExist == nil {
			continue
		}

		for _, out := range outsByScript[s] {
			newOut := &accountOutput{
				rawOutput: *out,
				AccountID: cp.AccountID,
				keyIndex:  cp.KeyIndex,
				change:    cp.Change,
			}
			result = append(result, newOut)
		}
	}

	return result
}

// upsertConfirmedAccountOutputs records the account data for confirmed utxos.
// If the account utxo already exists (because it's from a local tx), the
// block confirmation data will in the row will be updated.
func (m *Manager) upsertConfirmedAccountOutputs(outs []*accountOutput, block *legacy.Block, batch *db.Batch) error {
	var u *UTXO

	for _, out := range outs {
		u = &UTXO{OutputID: out.OutputID.Bytes(),
			AssetID:      out.AssetId.Bytes(),
			Amount:       out.Amount,
			AccountID:    out.AccountID,
			ProgramIndex: out.keyIndex,
			Program:      out.ControlProgram,
			SourceID:     out.sourceID.Bytes(),
			SourcePos:    out.sourcePos,
			RefData:      out.refData.Bytes(),
			Change:       out.change}

		rawUTXO, err := json.Marshal(u)
		if err != nil {
			return errors.Wrap(err, "failed marshal accountutxo")
		}

		(*batch).Set(accountUTXOKey(string(u.OutputID)), rawUTXO)
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/LeKovr/dbrpc/workman"
	"github.com/LeKovr/go-base/logger"
)

// -----------------------------------------------------------------------------

// ArgDef holds function argument attributes
type ArgDef struct {
	ID        int32
	Name      string
	Type      string
	Default   *string
	AllowNull bool
}

// FuncArgDef holds set of function argument attributes
type FuncArgDef []ArgDef

// -----------------------------------------------------------------------------

// FunctionDef creates a job for fetching of function argument definition
func FunctionDef(cfg *AplFlags, log *logger.Log, jc chan workman.Job, method string) (FuncArgDef, interface{}) {

	key := []string{cfg.ArgDefFunc, cfg.Schema + "." + method}

	payload, _ := json.Marshal(key)
	respChannel := make(chan workman.Result)

	work := workman.Job{Payload: string(payload), Result: respChannel}

	// Push the work onto the queue.
	jc <- work

	resp := <-respChannel
	log.Printf("Got def: %s", resp.Result)
	if !resp.Success {
		return nil, resp.Error
	}

	var res FuncArgDef
	err := json.Unmarshal(*resp.Result, &res)
	if err != nil {
		return nil, err
	}
	return res, nil
}

// -----------------------------------------------------------------------------

// FunctionResult creates a job for fetching of function result
func FunctionResult(jc chan workman.Job, payload string) workman.Result {

	respChannel := make(chan workman.Result)
	// let's create a job with the payload
	work := workman.Job{Payload: payload, Result: respChannel}

	// Push the work onto the queue.
	jc <- work

	resp := <-respChannel
	return resp
}

// -----------------------------------------------------------------------------

func httpHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		log.Printf("Request method: %s", r.Method)
		if r.Method == "GET" {
			getContextHandler(cfg, log, jc, w, r, true)
		} else if r.Method == "HEAD" {
			getContextHandler(cfg, log, jc, w, r, false) // Like get but without data
		} else if r.Method == "POST" && r.URL.Path == "/" {
			postContextHandler(cfg, log, jc, w, r)
		} else if r.Method == "POST" {
			postgrestContextHandler(cfg, log, jc, w, r)
		} else if r.Method == "OPTIONS" {
			optionsContextHandler(cfg, log, jc, w, r)
		} else {
			e := fmt.Sprintf("Unsupported request method: %s", r.Method)
			log.Warn(e)
			http.Error(w, e, http.StatusNotImplemented)
		}
	}
}

// -----------------------------------------------------------------------------

func getContextHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job, w http.ResponseWriter, r *http.Request, reply bool) {

	//	method := strings.TrimPrefix(strings.TrimPrefix(r.URL.Path, cfg.Prefix), "/")
	method := strings.TrimPrefix(r.URL.Path, cfg.Prefix)
	log.Printf("GotREquest %s (%s)", method, r.URL.Path)

	argDef, errd := FunctionDef(cfg, log, jc, method)
	if errd != nil {
		log.Warnf("Method %s load def error: %s", method, errd)
		http.NotFound(w, r)
		return
	}

	key := []*string{&method}
	r.ParseForm()

	f404 := []string{}
	for _, a := range argDef {
		v := r.Form[a.Name]

		if len(v) == 0 {
			if !a.AllowNull && a.Default == nil {
				f404 = append(f404, a.Name)
			} else if a.Default != nil {
				log.Debugf("Arg: %s use default", a.Name)
				break // use defaults
			}
			key = append(key, nil) // TODO: nil does not replaced with default
		} else {
			key = append(key, &v[0]) // TODO: array support
		}

		log.Debugf("Arg: %+v (%d)", v, len(f404))
	}
	var result workman.Result
	if len(f404) > 0 {
		result = workman.Result{Success: false, Error: fmt.Sprintf("Required parameter(s) %+v not found", f404)}
	} else {
		payload, _ := json.Marshal(key)
		log.Infof("Args: %s", string(payload))
		result = FunctionResult(jc, string(payload))
	}

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if reply {
		out, err := json.Marshal(result)
		if err != nil {
			log.Println("Marshall error: ", err)
			e := workman.Result{Success: false, Error: err.Error()}
			out, _ = json.Marshal(e)
		}
		w.Write(out)
		w.Write([]byte("\n"))
	}
}

// -----------------------------------------------------------------------------

func optionsContextHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job, w http.ResponseWriter, r *http.Request) {

	origin := r.Header.Get("Origin")
	var host string
	if origin != "" && len(cfg.Hosts) > 0 { // lookup if host is allowed
		for _, h := range cfg.Hosts {
			if origin == h {
				host = h
				break
			}
		}
	} else {
		host = origin
	}
	if origin != "" && host == "" {
		log.Warningf("Unregistered request source: %s", origin)
		http.Error(w, "Origin not registered", http.StatusForbidden)
		return
	}
	w.Header().Set("Access-Control-Allow-Origin", host)
	w.Header().Set("Access-Control-Allow-Headers", "origin, content-type, accept")
	w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
	w.WriteHeader(http.StatusOK)
	return
}

func getRaw(data interface{}) *json.RawMessage {
	j, _ := json.Marshal(data)
	raw := json.RawMessage(j)
	return &raw
}

// -----------------------------------------------------------------------------

// postContextHandler serve JSON-RPC envelope
func postContextHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job, w http.ResponseWriter, r *http.Request) {

	data, _ := ioutil.ReadAll(r.Body)
	req := serverRequest{}
	err := json.Unmarshal(data, &req)
	if err != nil {
		e := fmt.Sprintf("json parse error: %s", err)
		log.Warn(e)
		http.Error(w, e, http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	resultRPC := serverResponse{ID: req.ID, Version: req.Version}

	argDef, errd := FunctionDef(cfg, log, jc, req.Method)
	if errd != nil {
		log.Warnf("Method %s load def error: %s", req.Method, errd)
		resultRPC.Error = getRaw(respRPCError{Code: -32601, Message: errd.(string)})
	} else {
		// Load args
		key := []string{req.Method}
		r.ParseForm()
		f404 := []string{}
		for _, a := range argDef {
			v, ok := req.Params[a.Name]
			if !ok {
				if !a.AllowNull && a.Default == nil {
					f404 = append(f404, a.Name)
				}
				key = append(key, "") // TODO: nil
			} else {
				key = append(key, v.(string))
			}
		}

		if len(f404) > 0 {
			resultRPC.Error = getRaw(respRPCError{Code: -32602, Message: "Required parameter(s) not found", Data: getRaw(f404)})
		} else {
			payload, _ := json.Marshal(key)
			res := FunctionResult(jc, string(payload))
			if res.Success {
				resultRPC.Result = res.Result
			} else {
				resultRPC.Error = getRaw(respRPCError{Code: -32603, Message: "Internal Error", Data: getRaw(res.Error)})
			}
		}

	}

	out, err := json.Marshal(resultRPC)
	if err != nil {
		log.Println("Marshall error: ", err)
		resultRPC.Result = nil
		resultRPC.Error = getRaw(respRPCError{Code: -32603, Message: "Internal Error", Data: getRaw(err.Error())})

		out, _ = json.Marshal(resultRPC)
	}
	log.Debugf("JSON Resp: %s", string(out))
	w.Write(out)
	w.Write([]byte("\n"))
}

// -----------------------------------------------------------------------------

// postgrestContextHandler serve JSON-RPC envelope
// 404 when method not found
func postgrestContextHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job, w http.ResponseWriter, r *http.Request) {

	method := strings.TrimPrefix(r.URL.Path, cfg.Prefix)
	log.Debugf("postgrest call for %s", method)

	argDef, errd := FunctionDef(cfg, log, jc, method)
	if errd != nil {
		log.Warnf("Method %s load def error: %s", method, errd)
		http.NotFound(w, r)
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	resultStatus := http.StatusOK

	req := map[string]string{}
	var resultRPC interface{}

	data, _ := ioutil.ReadAll(r.Body)
	err := json.Unmarshal(data, &req)
	if err != nil {
		e := fmt.Sprintf("json parse error: %s", err)
		log.Warn(e)
		resultRPC = respPGTError{Message: "Cannot parse request payload", Details: e}
		resultStatus = http.StatusBadRequest
	} else {

		// Load args
		key := []string{method}
		f404 := []string{}
		for _, a := range argDef {
			v, ok := req[a.Name]
			if !ok {
				if !a.AllowNull && a.Default == nil {
					f404 = append(f404, a.Name)
				}
				key = append(key, "") // TODO: nil
			} else {
				key = append(key, v)
			}
		}

		if len(f404) > 0 {
			resultRPC = respPGTError{Code: "42883", Message: "Required parameter(s) not found", Details: strings.Join(f404, ", ")}
			resultStatus = http.StatusBadRequest
		} else {
			payload, _ := json.Marshal(key)
			res := FunctionResult(jc, string(payload))
			if res.Success {
				resultRPC = res.Result
			} else {
				resultRPC = respPGTError{Message: "Method call error", Details: res.Error.(string)}
				resultStatus = http.StatusBadRequest // TODO: ?
			}
		}

	}

	out, err := json.Marshal(resultRPC)
	if err != nil {
		log.Println("Marshall error: ", err)
		resultRPC = respPGTError{Message: "Method result marshall error", Details: err.Error()}
		resultStatus = http.StatusBadRequest // TODO: ?
		out, _ = json.Marshal(resultRPC)
	}
	w.WriteHeader(resultStatus)
	log.Debugf("JSON Resp: %s", string(out))
	w.Write(out)
	w.Write([]byte("\n"))
}

// JSON-RPC v2.0 structures

type serverRequest struct {
	Method  string                 `json:"method"`
	Version string                 `json:"jsonrpc"`
	ID      uint64                 `json:"id"`
	Params  map[string]interface{} `json:"params"`
}

type serverResponse struct {
	ID      uint64           `json:"id"`
	Version string           `json:"jsonrpc"`
	Result  *json.RawMessage `json:"result,omitempty"`
	Error   *json.RawMessage `json:"error,omitempty"`
}

type respRPCError struct {
	Code    int              `json:"code"`
	Message string           `json:"message"`
	Data    *json.RawMessage `json:"data,omitempty"`
}
type respPGTError struct {
	Message string `json:"message"`
	Code    string `json:"code,omitempty"`
	Details string `json:"details,omitempty"`
}
<commit_msg>debug print improved, added metrics & array support<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/LeKovr/dbrpc/workman"
	"github.com/LeKovr/go-base/logger"
)

// -----------------------------------------------------------------------------

// ArgDef holds function argument attributes
type ArgDef struct {
	ID        int32
	Name      string
	Type      string
	Default   *string
	AllowNull bool
}

// FuncArgDef holds set of function argument attributes
type FuncArgDef []ArgDef

// JSON-RPC v2.0 structures
type reqParams map[string]interface{}

type serverRequest struct {
	Method  string    `json:"method"`
	Version string    `json:"jsonrpc"`
	ID      uint64    `json:"id"`
	Params  reqParams `json:"params"`
}

type serverResponse struct {
	ID      uint64           `json:"id"`
	Version string           `json:"jsonrpc"`
	Result  *json.RawMessage `json:"result,omitempty"`
	Error   *json.RawMessage `json:"error,omitempty"`
}

type respRPCError struct {
	Code    int              `json:"code"`
	Message string           `json:"message"`
	Data    *json.RawMessage `json:"data,omitempty"`
}
type respPGTError struct {
	Message string `json:"message"`
	Code    string `json:"code,omitempty"`
	Details string `json:"details,omitempty"`
}

// -----------------------------------------------------------------------------

// FunctionDef creates a job for fetching of function argument definition
func FunctionDef(cfg *AplFlags, log *logger.Log, jc chan workman.Job, method string) (FuncArgDef, interface{}) {

	key := []string{cfg.ArgDefFunc, cfg.Schema + "." + method}

	payload, _ := json.Marshal(key)
	respChannel := make(chan workman.Result)

	work := workman.Job{Payload: string(payload), Result: respChannel}

	// Push the work onto the queue.
	jc <- work

	resp := <-respChannel
	log.Debugf("Got def: %s", resp.Result)
	if !resp.Success {
		return nil, resp.Error
	}

	var res FuncArgDef
	err := json.Unmarshal(*resp.Result, &res)
	if err != nil {
		return nil, err
	}
	return res, nil
}

// -----------------------------------------------------------------------------

// FunctionResult creates a job for fetching of function result
func FunctionResult(jc chan workman.Job, payload string) workman.Result {

	respChannel := make(chan workman.Result)
	// let's create a job with the payload
	work := workman.Job{Payload: payload, Result: respChannel}

	// Push the work onto the queue.
	jc <- work

	resp := <-respChannel
	return resp
}

// -----------------------------------------------------------------------------

func httpHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		log.Debugf("Request method: %s", r.Method)
		if r.Method == "GET" {
			getContextHandler(cfg, log, jc, w, r, true)
		} else if r.Method == "HEAD" {
			getContextHandler(cfg, log, jc, w, r, false) // Like get but without data
		} else if r.Method == "POST" && r.URL.Path == cfg.Prefix {
			postContextHandler(cfg, log, jc, w, r)
		} else if r.Method == "POST" {
			postgrestContextHandler(cfg, log, jc, w, r)
		} else if r.Method == "OPTIONS" {
			optionsContextHandler(cfg, log, jc, w, r)
		} else {
			e := fmt.Sprintf("Unsupported request method: %s", r.Method)
			log.Warn(e)
			http.Error(w, e, http.StatusNotImplemented)
		}
	}
}

// -----------------------------------------------------------------------------

func getContextHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job, w http.ResponseWriter, r *http.Request, reply bool) {
	start := time.Now()

	method := strings.TrimPrefix(r.URL.Path, cfg.Prefix)

	argDef, errd := FunctionDef(cfg, log, jc, method)
	if errd != nil {
		// Warning was when fetched from db
		log.Infof("Method %s load def error: %s", method, errd)
		http.NotFound(w, r)
		return
	}

	key := []*string{&method}
	r.ParseForm()

	f404 := []string{}
	for _, a := range argDef {
		v := r.Form[a.Name]

		if len(v) == 0 {
			if !a.AllowNull && a.Default == nil {
				f404 = append(f404, a.Name)
			} else if a.Default != nil {
				log.Debugf("Arg: %s use default", a.Name)
				break // use defaults
			}
			key = append(key, nil) // TODO: nil does not replaced with default
		} else if strings.HasSuffix(a.Type, "[]") {
			// convert array into string
			// TODO: escape ","
			s := "{" + strings.Join(v, ",") + "}"
			key = append(key, &s)
		} else {
			key = append(key, &v[0])
		}

		log.Debugf("Arg: %+v (%d)", v, len(f404))
	}
	var result workman.Result
	if len(f404) > 0 {
		result = workman.Result{Success: false, Error: fmt.Sprintf("Required parameter(s) %+v not found", f404)}
	} else {
		payload, _ := json.Marshal(key)
		log.Debugf("Args: %s", string(payload))
		result = FunctionResult(jc, string(payload))
	}

	if reply {
		out, err := json.Marshal(result)
		if err != nil {
			log.Warnf("Marshall error: %+v", err)
			e := workman.Result{Success: false, Error: err.Error()}
			out, _ = json.Marshal(e)
		}
		setMetric(w, start, http.StatusOK)
		w.Write(out)
		//w.Write([]byte("\n"))
	} else {
		w.WriteHeader(http.StatusOK)
	}
}

// -----------------------------------------------------------------------------

func setMetric(w http.ResponseWriter, start time.Time, status int) {
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Header().Set("X-Elapsed", fmt.Sprint(time.Since(start)))
	w.WriteHeader(http.StatusOK)
}

// -----------------------------------------------------------------------------

func optionsContextHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job, w http.ResponseWriter, r *http.Request) {

	origin := r.Header.Get("Origin")
	var host string
	if origin != "" && len(cfg.Hosts) > 0 { // lookup if host is allowed
		for _, h := range cfg.Hosts {
			if origin == h {
				host = h
				break
			}
		}
	} else {
		host = origin
	}
	if origin != "" && host == "" {
		log.Warningf("Unregistered request source: %s", origin)
		http.Error(w, "Origin not registered", http.StatusForbidden)
		return
	}
	w.Header().Set("Access-Control-Allow-Origin", host)
	w.Header().Set("Access-Control-Allow-Headers", "origin, content-type, accept")
	w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
	w.WriteHeader(http.StatusOK)
	return
}

func getRaw(data interface{}) *json.RawMessage {
	j, _ := json.Marshal(data)
	raw := json.RawMessage(j)
	return &raw
}

// -----------------------------------------------------------------------------

// postContextHandler serve JSON-RPC envelope
func postContextHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job, w http.ResponseWriter, r *http.Request) {

	start := time.Now()

	data, _ := ioutil.ReadAll(r.Body)
	req := serverRequest{}
	err := json.Unmarshal(data, &req)
	if err != nil {
		e := fmt.Sprintf("json parse error: %s", err)
		log.Warn(e)
		http.Error(w, e, http.StatusBadRequest)
		return
	}

	resultRPC := serverResponse{ID: req.ID, Version: req.Version}

	argDef, errd := FunctionDef(cfg, log, jc, req.Method)
	if errd != nil {
		log.Warnf("Method %s load def error: %s", req.Method, errd)
		resultRPC.Error = getRaw(respRPCError{Code: -32601, Message: errd.(string)})
	} else {
		// Load args
		r.ParseForm()
		key, f404 := fetchArgs(argDef, req.Params, req.Method)
		if len(f404) > 0 {
			resultRPC.Error = getRaw(respRPCError{Code: -32602, Message: "Required parameter(s) not found", Data: getRaw(f404)})
		} else {
			payload, _ := json.Marshal(key)
			log.Debugf("Args: %s", string(payload))
			res := FunctionResult(jc, string(payload))
			if res.Success {
				resultRPC.Result = res.Result
			} else {
				resultRPC.Error = getRaw(respRPCError{Code: -32603, Message: "Internal Error", Data: getRaw(res.Error)})
			}
		}
	}

	out, err := json.Marshal(resultRPC)
	if err != nil {
		log.Warnf("Marshall error: %+v", err)
		resultRPC.Result = nil
		resultRPC.Error = getRaw(respRPCError{Code: -32603, Message: "Internal Error", Data: getRaw(err.Error())})

		out, _ = json.Marshal(resultRPC)
	}
	log.Debugf("JSON Resp: %s", string(out))

	setMetric(w, start, http.StatusOK)
	w.Write(out)
	//w.Write([]byte("\n"))
}

// -----------------------------------------------------------------------------

// postgrestContextHandler serve JSON-RPC envelope
// 404 when method not found
func postgrestContextHandler(cfg *AplFlags, log *logger.Log, jc chan workman.Job, w http.ResponseWriter, r *http.Request) {

	start := time.Now()

	method := strings.TrimPrefix(r.URL.Path, cfg.Prefix)
	log.Debugf("postgrest call for %s", method)

	argDef, errd := FunctionDef(cfg, log, jc, method)
	if errd != nil {
		log.Warnf("Method %s load def error: %s", method, errd)
		http.NotFound(w, r)
		return
	}
	resultStatus := http.StatusOK

	req := reqParams{}
	var resultRPC interface{}

	data, _ := ioutil.ReadAll(r.Body)
	err := json.Unmarshal(data, &req)

	if err != nil {
		e := fmt.Sprintf("json parse error: %s", err)
		log.Warn(e)
		resultRPC = respPGTError{Message: "Cannot parse request payload", Details: e}
		resultStatus = http.StatusBadRequest
	} else {
		// Load args
		key, f404 := fetchArgs(argDef, req, method)
		if len(f404) > 0 {
			resultRPC = respPGTError{Code: "42883", Message: "Required parameter(s) not found", Details: strings.Join(f404, ", ")}
			resultStatus = http.StatusBadRequest
		} else {
			payload, _ := json.Marshal(key)
			log.Debugf("Args: %s", string(payload))
			res := FunctionResult(jc, string(payload))
			if res.Success {
				resultRPC = res.Result
			} else {
				resultRPC = respPGTError{Message: "Method call error", Details: res.Error.(string)}
				resultStatus = http.StatusBadRequest // TODO: ?
			}
		}
	}

	out, err := json.Marshal(resultRPC)
	if err != nil {
		log.Warnf("Marshall error: %+v", err)
		resultRPC = respPGTError{Message: "Method result marshall error", Details: err.Error()}
		resultStatus = http.StatusBadRequest // TODO: ?
		out, _ = json.Marshal(resultRPC)
	}
	setMetric(w, start, resultStatus)
	log.Debugf("JSON Resp: %s", string(out))
	w.Write(out)
	//w.Write([]byte("\n"))
}

func fetchArgs(argDef FuncArgDef, req reqParams, method string) ([]*string, []string) {

	key := []*string{&method}
	f404 := []string{}

	for _, a := range argDef {
		v, ok := req[a.Name]
		if !ok {
			if !a.AllowNull && a.Default == nil {
				f404 = append(f404, a.Name)
			} else if a.Default != nil {
				//log.Debugf("Arg: %s use default", a.Name)
				break // use defaults
			}
			key = append(key, nil) // TODO: nil does not replaced with default
		} else if strings.HasSuffix(a.Type, "[]") {
			// wait slice
			s := reflect.ValueOf(v)
			if s.Kind() != reflect.Slice {
				// string or {string}
				vs := v.(string)

				// // convert scalar to postgres array
				// asArray = regexp.MustCompile(`^\{.+\}$`)
				// if !asArray.MatchString(vs) {
				// 	vs = "{" + vs + "}"
				// }
				key = append(key, &vs)
			} else {
				// slice
				ret := make([]string, s.Len())

				for i := 0; i < s.Len(); i++ {
					ret[i] = s.Index(i).Interface().(string)
					//	log.Printf("====== %+v", ret[i])
				}
				// convert array into string
				// TODO: escape ","
				ss := "{" + strings.Join(ret, ",") + "}"
				key = append(key, &ss)
			}
		} else {
			s := v.(string)
			key = append(key, &s)
		}

		//log.Debugf("Arg: %+v (%d)", v, len(f404))
	}
	return key, f404
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"github.com/golang/groupcache"
	"github.com/gorilla/mux"
	"io"
	"math/rand"
	"net/http"
	"time"
	"vip/fetch"
)

type UploadResponse struct {
	Url string `json:"url"`
}

type verifyAuth func(http.ResponseWriter, *http.Request)

func (h verifyAuth) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	auth := r.Header.Get("X-Vip-Token")
	if auth != authToken {
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	h(w, r)
}

func fileKey(bucket string) string {
	seed := rand.New(rand.NewSource(time.Now().UnixNano()))
	key := fmt.Sprintf("%d-%s-%d", seed.Int63(), bucket, time.Now().UnixNano())

	hash := md5.New()
	io.WriteString(hash, key)
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func handleImageRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(http.StatusMethodNotAllowed)
	}

	w.Header().Set("Cache-Control", "public, max-age=31536000")

	// Client is checking for a cached URI, assume it is valid
	// and return a 304
	if r.Header.Get("If-Modified-Since") != "" {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	gc := fetch.RequestContext(r)

	var data []byte
	err := cache.Get(gc, gc.CacheKey(), groupcache.AllocatingByteSliceSink(&data))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	w.Header().Set("Content-Type", http.DetectContentType(data))
	http.ServeContent(w, r, gc.ImageId, time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC), bytes.NewReader(data))

	return
}

func handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
	}

	vars := mux.Vars(r)
	bucket := vars["bucket_id"]

	// Set a hard 5mb limit on files
	if r.ContentLength > 6<<20 {
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		return
	}

	key := fileKey(bucket)
	err := storage.PutReader(bucket, key, r.Body,
		r.ContentLength, r.Header.Get("Content-Type"))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusCreated)
	json.NewEncoder(w).Encode(UploadResponse{
		Url: fmt.Sprintf("%s/%s", bucket, key),
	})

	return
}

func handlePing(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, "pong")

	return
}
<commit_msg>Remove superfluous returns<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"github.com/golang/groupcache"
	"github.com/gorilla/mux"
	"io"
	"math/rand"
	"net/http"
	"time"
	"vip/fetch"
)

type UploadResponse struct {
	Url string `json:"url"`
}

type verifyAuth func(http.ResponseWriter, *http.Request)

func (h verifyAuth) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	auth := r.Header.Get("X-Vip-Token")
	if auth != authToken {
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	h(w, r)
}

func fileKey(bucket string) string {
	seed := rand.New(rand.NewSource(time.Now().UnixNano()))
	key := fmt.Sprintf("%d-%s-%d", seed.Int63(), bucket, time.Now().UnixNano())

	hash := md5.New()
	io.WriteString(hash, key)
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func handleImageRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(http.StatusMethodNotAllowed)
	}

	w.Header().Set("Cache-Control", "public, max-age=31536000")

	// Client is checking for a cached URI, assume it is valid
	// and return a 304
	if r.Header.Get("If-Modified-Since") != "" {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	gc := fetch.RequestContext(r)

	var data []byte
	err := cache.Get(gc, gc.CacheKey(), groupcache.AllocatingByteSliceSink(&data))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	w.Header().Set("Content-Type", http.DetectContentType(data))
	http.ServeContent(w, r, gc.ImageId, time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC), bytes.NewReader(data))
}

func handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
	}

	vars := mux.Vars(r)
	bucket := vars["bucket_id"]

	// Set a hard 5mb limit on files
	if r.ContentLength > 6<<20 {
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		return
	}

	key := fileKey(bucket)
	err := storage.PutReader(bucket, key, r.Body,
		r.ContentLength, r.Header.Get("Content-Type"))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusCreated)
	json.NewEncoder(w).Encode(UploadResponse{
		Url: fmt.Sprintf("%s/%s", bucket, key),
	})
}

func handlePing(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, "pong")
}
<|endoftext|>


<commit_before>package tori

import "bufio"
import "bytes"
import "fmt"
import "net/http"
import "github.com/shiroyuki/re"

type Handler struct {
    Route         *Route
    Response      *http.ResponseWriter
    Request       *http.Request
    Parameters    *re.MultipleResult
    BufferEnabled bool
    Status        uint16
    buffer        *bytes.Buffer
}

// Construct a handler.
//
// The constructed handler will use a buffer by default. This is to allow a
// middleware to work on the response content and minimize the I/O interuption.
//
// Disabling buffering will automatically prevent a middleware from doing
// post-processing on the response content for the corresponding request.
func NewHandler(
    route      *Route,
    response   *http.ResponseWriter,
    request    *http.Request,
    parameters *re.MultipleResult,
) *Handler {
    handler := &Handler{
        Route:         route,
        Response:      response,
        Request:       request,
        Parameters:    parameters,
        BufferEnabled: true,
    }

    handler.SetStatus(http.StatusOK)

    return handler
}

// Set the HTTP status code for the response.
func (self *Handler) SetStatus(statusCode uint16) {
    self.Status = statusCode
}

// Get the request header.
func (self *Handler) GetHeader(key string) string {
    return self.Request.Header.Get(key)
}

// Add the response header.
func (self *Handler) AddHeader(key string, value string) {
    (*self.Response).Header().Add(key, value)
}

// Set the response header.
func (self *Handler) SetHeader(key string, value string) {
    (*self.Response).Header().Set(key, value)
}

func (self *Handler) SetContentType(contentType string) {
    self.SetHeader("Content-Type", contentType)
}

func (self *Handler) SetContentLength(contentLength int) {
    self.SetHeader("Content-Length", fmt.Sprintf("%d", contentLength))
}

func (self *Handler) SetContentEncoding(encoding string) {
    self.SetHeader("Content-Encoding", encoding)
}

// Disable buffering.
func (self *Handler) DisableBuffering() {
    self.BufferEnabled = false
}

func (self *Handler) Write(content string) {
    self.WriteByte([]byte(content))
}

func (self *Handler) WriteByte(content []byte) {
    if self.BufferEnabled {
        if self.buffer == nil {
            self.buffer = new(bytes.Buffer)
        }

        w := bufio.NewWriter(self.buffer)

        defer w.Flush()

        w.Write(content)

        return
    }

    (*self.Response).Write(content)
}

func (self *Handler) Content() []byte {
    return self.buffer.Bytes()
}
<commit_msg>Quick updates<commit_after>package tori

import "bufio"
import "bytes"
import "fmt"
import "net/http"
import "github.com/shiroyuki/re"

type Handler struct {
    Route         *Route
    Response      *http.ResponseWriter
    Request       *http.Request
    Parameters    *re.MultipleResult
    BufferEnabled bool
    Status        uint16
    buffer        *bytes.Buffer
}

// Construct a handler.
//
// The constructed handler will use a buffer by default. This is to allow a
// middleware to work on the response content and minimize the I/O interuption.
//
// Disabling buffering will automatically prevent a middleware from doing
// post-processing on the response content for the corresponding request.
func NewHandler(
    route      *Route,
    response   *http.ResponseWriter,
    request    *http.Request,
    parameters *re.MultipleResult,
) *Handler {
    handler := &Handler{
        Route:         route,
        Response:      response,
        Request:       request,
        Parameters:    parameters,
        BufferEnabled: true,
    }

    handler.SetStatus(http.StatusOK)

    return handler
}

// Set the HTTP status code for the response.
func (self *Handler) SetStatus(statusCode uint16) {
    self.Status = statusCode
}

// Get the request header.
func (self *Handler) GetHeader(key string) string {
    return self.Request.Header.Get(key)
}

// Retrieves the parameter at the match index of the request path.
//
// For example, if the non-reversible pattern is defined as /path/(\d+)/([a-z]+)
// and the request path is /path/123/def, there will be two indices where the
// values are "123" and "def" for index 1 and 2 respectively.
func (self *Handler) Index(key int) string {
    value := self.Parameters.Index(key)

    return *value
}

// Retrieves the parameter defined by the key of the request path.
//
// For example, if the reversible pattern is defined as /path/<k1>/<h2> and
// the request path is /path/abc/def, there will be two keys, "k1" and "k2".
func (self *Handler) Key(key string) []string {
    value := self.Parameters.Key(key)

    return *value
}

// Add the response header.
func (self *Handler) AddHeader(key string, value string) {
    (*self.Response).Header().Add(key, value)
}

// Set the response header.
func (self *Handler) SetHeader(key string, value string) {
    (*self.Response).Header().Set(key, value)
}

func (self *Handler) SetContentType(contentType string) {
    self.SetHeader("Content-Type", contentType)
}

func (self *Handler) SetContentLength(contentLength int) {
    self.SetHeader("Content-Length", fmt.Sprintf("%d", contentLength))
}

func (self *Handler) SetContentEncoding(encoding string) {
    self.SetHeader("Content-Encoding", encoding)
}

// Disable buffering.
func (self *Handler) DisableBuffering() {
    self.BufferEnabled = false
}

func (self *Handler) Write(content string) {
    self.WriteByte([]byte(content))
}

func (self *Handler) WriteByte(content []byte) {
    if self.BufferEnabled {
        if self.buffer == nil {
            self.buffer = new(bytes.Buffer)
        }

        w := bufio.NewWriter(self.buffer)

        defer w.Flush()

        w.Write(content)

        return
    }

    (*self.Response).Write(content)
}

func (self *Handler) Content() []byte {
    return self.buffer.Bytes()
}
<|endoftext|>


<commit_before>package uploader

import (
	"net/http"
)

func UploadHandler(uploader *Uploader, uploadField, filenameField string) http.HandlerFunc {

	return func(res http.ResponseWriter, req *http.Request) {
		req.ParseForm()

		switch req.Method {
		case "GET":
			filename := req.FormValue(filenameField)
			if filename == "" {
				res.WriteHeader(http.StatusBadRequest)
				return
			}
			imageData, err := uploader.Get(filename)
			if err != nil {
				res.WriteHeader(http.StatusInternalServerError)
				return
			}
			res.WriteHeader(http.StatusOK)
			res.Write(imageData)
			return
		case "POST":
			//file, fileinfo, err := req.FormFile(uploadField)

		case "PUT":

		case "DELETE":
			filename := req.FormValue(filenameField)
			if filename == "" {
				res.WriteHeader(http.StatusBadRequest)
				return
			}
			err := uploader.Delete(filename)
			if err != nil {
				res.WriteHeader(http.StatusInternalServerError)
				return
			}
			res.WriteHeader(http.StatusOK)
			res.Write([]byte("{\"status\": \"OK\"}"))
		default:
			res.WriteHeader(http.StatusMethodNotAllowed)
		}

	}
}
<commit_msg>Implement PUT/POST<commit_after>package uploader

import (
	"fmt"
	"io/ioutil"
	"net/http"
)

func UploadHandler(uploader *Uploader, uploadField, filenameField string) http.HandlerFunc {

	return func(res http.ResponseWriter, req *http.Request) {
		req.ParseForm()

		switch req.Method {
		case "GET":
			filename := req.FormValue(filenameField)
			if filename == "" {
				res.WriteHeader(http.StatusBadRequest)
				return
			}
			imageData, err := uploader.Get(filename)
			if err != nil {
				if !uploader.Has(filename) {
					res.WriteHeader(http.StatusNotFound)
					return
				}
				res.WriteHeader(http.StatusInternalServerError)
				return
			}
			res.WriteHeader(http.StatusOK)
			res.Write(imageData)
			return
		case "PUT":
			// same for now
			fallthrough
		case "POST":
			file, fileinfo, err := req.FormFile(uploadField)
			if err != nil {
				res.WriteHeader(http.StatusBadRequest)
				return
			}
			fmt.Println(file, fileinfo)
			imageData, err := ioutil.ReadAll(file)
			if err != nil {
				res.WriteHeader(http.StatusInternalServerError)
				return
			}

			filename, err := uploader.Store(imageData)
			if err != nil {
				res.WriteHeader(http.StatusInternalServerError)
				return
			}
			res.WriteHeader(http.StatusOK)
			res.Write([]byte(fmt.Sprintf("{\"status\": \"OK\", \"filename\": \"%s\"}", filename)))
			return
		case "DELETE":
			filename := req.FormValue(filenameField)
			if filename == "" {
				res.WriteHeader(http.StatusBadRequest)
				return
			}
			err := uploader.Delete(filename)
			if err != nil {
				if !uploader.Has(filename) {
					res.WriteHeader(http.StatusNotFound)
					return
				}
				res.WriteHeader(http.StatusInternalServerError)
				return
			}
			res.WriteHeader(http.StatusOK)
			res.Write([]byte("{\"status\": \"OK\"}"))
		default:
			res.WriteHeader(http.StatusMethodNotAllowed)
		}

	}
}
<|endoftext|>


<commit_before>package main

type Haproxy struct {
	Socket string
}

//func (h *Haproxy) ff() {
//
//}
<commit_msg>start laying out architecture for tweaking the config<commit_after>package main

type Haproxy struct {
	Socket string
}

func (h *Haproxy) GracefulRestart() {

}

//func (h *Haproxy) AddFrontend() {

//}

func (h *Haproxy) WriteConfig() {

}

func (h *Haproxy) ReadConfig() {

}
<|endoftext|>


<commit_before>package colog

import (
	"bytes"
	"encoding/json"
	"sort"
)

type HashSet map[Hash]struct{}

func (s HashSet) Unset(h Hash) {
	delete(s, h)
}

func (s HashSet) Set(h Hash) {
	s[h] = struct{}{}
}

func (s HashSet) IsSet(h Hash) bool {
	_, ok := s[h]
	return ok
}

func (s HashSet) Copy() HashSet {
	s_ := make(HashSet)

	for k, v := range s {
		s_[k] = v
	}

	return s_
}

func (s HashSet) Sorted() []Hash {
	hs := make([]Hash, 0, len(s))
	ss := make([]string, 0, len(s))

	for h := range s {
		ss = append(ss, string(h))
	}

	// TODO find out if this sorts good enough
	sort.Strings(ss)

	for _, s := range ss {
		hs = append(hs, Hash(s))
	}

	return hs
}

func (s HashSet) MarshalJSON() ([]byte, error) {
	return json.Marshal(s.Sorted())
}

func (s HashSet) UnmarshalJSON(in []byte) error {
	var hs []Hash

	err := json.Unmarshal(in, &hs)
	if err != nil {
		return err
	}

	for _, h := range hs {
		s.Set(h)
	}

	return nil
}

func (s HashSet) String() string {

	hs := s.Sorted()

	buf := bytes.NewBufferString("{ ")

	if len(hs) > 0 {
		buf.WriteString(hs[0].String())

		hs = hs[1:]

		for _, h := range hs {
			buf.WriteString(", " + h.String())
		}

	}

	buf.WriteString(" }")

	return buf.String()
}
<commit_msg>slight change<commit_after>package colog

import (
	"bytes"
	"encoding/json"
	"sort"
)

type HashSet map[Hash]struct{}

func (s HashSet) Unset(h Hash) {
	delete(s, h)
}

func (s HashSet) Set(h Hash) {
	s[h] = struct{}{}
}

func (s HashSet) IsSet(h Hash) bool {
	_, ok := s[h]
	return ok
}

func (s HashSet) Copy() HashSet {
	s_ := make(HashSet)

	for k := range s {
		s_.Set(k)
	}

	return s_
}

func (s HashSet) Sorted() []Hash {
	hs := make([]Hash, 0, len(s))
	ss := make([]string, 0, len(s))

	for h := range s {
		ss = append(ss, string(h))
	}

	// TODO find out if this sorts good enough
	sort.Strings(ss)

	for _, s := range ss {
		hs = append(hs, Hash(s))
	}

	return hs
}

func (s HashSet) MarshalJSON() ([]byte, error) {
	return json.Marshal(s.Sorted())
}

func (s HashSet) UnmarshalJSON(in []byte) error {
	var hs []Hash

	err := json.Unmarshal(in, &hs)
	if err != nil {
		return err
	}

	for _, h := range hs {
		s.Set(h)
	}

	return nil
}

func (s HashSet) String() string {

	hs := s.Sorted()

	buf := bytes.NewBufferString("{ ")

	if len(hs) > 0 {
		buf.WriteString(hs[0].String())

		hs = hs[1:]

		for _, h := range hs {
			buf.WriteString(", " + h.String())
		}

	}

	buf.WriteString(" }")

	return buf.String()
}
<|endoftext|>