<commit_before>package main

import (
	"context"
	"fmt"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/cio"
	"github.com/containerd/containerd/oci"
	metrics "github.com/docker/go-metrics"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/sirupsen/logrus"
)

var ct metrics.LabeledTimer

func init() {
	ns := metrics.NewNamespace("stress", "", nil)
	// if you want more fine grained metrics then you can drill down with the metrics in prom that
	// containerd is outputing
	ct = ns.NewLabeledTimer("run", "Run time of a full container during the test", "commit")
	metrics.Register(ns)
}

type worker struct {
	id       int
	wg       *sync.WaitGroup
	count    int
	failures int

	client *containerd.Client
	image  containerd.Image
	spec   *specs.Spec
	doExec bool
	commit string
}

func (w *worker) run(ctx, tctx context.Context) {
	defer func() {
		w.wg.Done()
		logrus.Infof("worker %d finished", w.id)
	}()
	for {
		select {
		case <-tctx.Done():
			return
		default:
		}

		w.count++
		id := w.getID()
		logrus.Debugf("starting container %s", id)
		start := time.Now()
		if err := w.runContainer(ctx, id); err != nil {
			if err != context.DeadlineExceeded ||
				!strings.Contains(err.Error(), context.DeadlineExceeded.Error()) {
				w.failures++
				logrus.WithError(err).Errorf("running container %s", id)

			}
			continue
		}
		// only log times are success so we don't scew the results from failures that go really fast
		ct.WithValues(w.commit).UpdateSince(start)
	}
}

func (w *worker) runContainer(ctx context.Context, id string) error {
	// fix up cgroups path for a default config
	w.spec.Linux.CgroupsPath = filepath.Join("/", "stress", id)
	c, err := w.client.NewContainer(ctx, id,
		containerd.WithNewSnapshot(id, w.image),
		containerd.WithSpec(w.spec, oci.WithUsername("games")),
	)
	if err != nil {
		return err
	}
	defer c.Delete(ctx, containerd.WithSnapshotCleanup)

	task, err := c.NewTask(ctx, cio.NullIO)
	if err != nil {
		return err
	}
	defer task.Delete(ctx, containerd.WithProcessKill)

	statusC, err := task.Wait(ctx)
	if err != nil {
		return err
	}
	if err := task.Start(ctx); err != nil {
		return err
	}
	if w.doExec {
		for i := 0; i < 256; i++ {
			if err := w.exec(ctx, i, task); err != nil {
				w.failures++
				logrus.WithError(err).Error("exec failure")
			}
		}
		if err := task.Kill(ctx, syscall.SIGKILL); err != nil {
			return err
		}
	}
	status := <-statusC
	_, _, err = status.Result()
	if err != nil {
		if err == context.DeadlineExceeded || err == context.Canceled {
			return nil
		}
		w.failures++
	}
	return nil
}

func (w *worker) exec(ctx context.Context, i int, t containerd.Task) error {
	pSpec := *w.spec.Process
	pSpec.Args = []string{"true"}
	process, err := t.Exec(ctx, strconv.Itoa(i), &pSpec, cio.NullIO)
	if err != nil {
		return err
	}
	defer process.Delete(ctx)
	status, err := process.Wait(ctx)
	if err != nil {
		return err
	}
	if err := process.Start(ctx); err != nil {
		return err
	}
	<-status
	return nil
}

func (w *worker) getID() string {
	return fmt.Sprintf("%d-%d", w.id, w.count)
}
<commit_msg>Add error metric for stress tests<commit_after>package main

import (
	"context"
	"fmt"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/cio"
	"github.com/containerd/containerd/oci"
	metrics "github.com/docker/go-metrics"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/sirupsen/logrus"
)

var (
	ct         metrics.LabeledTimer
	errCounter metrics.LabeledCounter
)

func init() {
	ns := metrics.NewNamespace("stress", "", nil)
	// if you want more fine grained metrics then you can drill down with the metrics in prom that
	// containerd is outputing
	ct = ns.NewLabeledTimer("run", "Run time of a full container during the test", "commit")
	errCounter = ns.NewLabeledCounter("errors", "Errors encountered running the stress tests", "err")
	metrics.Register(ns)
}

type worker struct {
	id       int
	wg       *sync.WaitGroup
	count    int
	failures int

	client *containerd.Client
	image  containerd.Image
	spec   *specs.Spec
	doExec bool
	commit string
}

func (w *worker) run(ctx, tctx context.Context) {
	defer func() {
		w.wg.Done()
		logrus.Infof("worker %d finished", w.id)
	}()
	for {
		select {
		case <-tctx.Done():
			return
		default:
		}

		w.count++
		id := w.getID()
		logrus.Debugf("starting container %s", id)
		start := time.Now()
		if err := w.runContainer(ctx, id); err != nil {
			if err != context.DeadlineExceeded ||
				!strings.Contains(err.Error(), context.DeadlineExceeded.Error()) {
				w.failures++
				logrus.WithError(err).Errorf("running container %s", id)
				errCounter.WithValues(err.Error()).Inc()

			}
			continue
		}
		// only log times are success so we don't scew the results from failures that go really fast
		ct.WithValues(w.commit).UpdateSince(start)
	}
}

func (w *worker) runContainer(ctx context.Context, id string) error {
	// fix up cgroups path for a default config
	w.spec.Linux.CgroupsPath = filepath.Join("/", "stress", id)
	c, err := w.client.NewContainer(ctx, id,
		containerd.WithNewSnapshot(id, w.image),
		containerd.WithSpec(w.spec, oci.WithUsername("games")),
	)
	if err != nil {
		return err
	}
	defer c.Delete(ctx, containerd.WithSnapshotCleanup)

	task, err := c.NewTask(ctx, cio.NullIO)
	if err != nil {
		return err
	}
	defer task.Delete(ctx, containerd.WithProcessKill)

	statusC, err := task.Wait(ctx)
	if err != nil {
		return err
	}
	if err := task.Start(ctx); err != nil {
		return err
	}
	if w.doExec {
		for i := 0; i < 256; i++ {
			if err := w.exec(ctx, i, task); err != nil {
				w.failures++
				logrus.WithError(err).Error("exec failure")
			}
		}
		if err := task.Kill(ctx, syscall.SIGKILL); err != nil {
			return err
		}
	}
	status := <-statusC
	_, _, err = status.Result()
	if err != nil {
		if err == context.DeadlineExceeded || err == context.Canceled {
			return nil
		}
		w.failures++
	}
	return nil
}

func (w *worker) exec(ctx context.Context, i int, t containerd.Task) error {
	pSpec := *w.spec.Process
	pSpec.Args = []string{"true"}
	process, err := t.Exec(ctx, strconv.Itoa(i), &pSpec, cio.NullIO)
	if err != nil {
		return err
	}
	defer process.Delete(ctx)
	status, err := process.Wait(ctx)
	if err != nil {
		return err
	}
	if err := process.Start(ctx); err != nil {
		return err
	}
	<-status
	return nil
}

func (w *worker) getID() string {
	return fmt.Sprintf("%d-%d", w.id, w.count)
}
<|endoftext|>


<commit_before>package collect

import (
	"bytes"
	"compress/gzip"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/slog"
)

func send() {
	for {
		qlock.Lock()
		if len(queue) > 0 {
			i := len(queue)
			if i > BatchSize {
				i = BatchSize
			}
			sending := queue[:i]
			queue = queue[i:]
			qlock.Unlock()
			slog.Infof("sending: %d, remaining: %d", len(sending), len(queue))
			sendBatch(sending)
		} else {
			qlock.Unlock()
			time.Sleep(time.Second)
		}
	}
}

func sendBatch(batch opentsdb.MultiDataPoint) {
	b, err := batch.Json()
	if err != nil {
		slog.Error(3, err)
		// bad JSON encoding, just give up
		return
	}
	var buf bytes.Buffer
	g := gzip.NewWriter(&buf)
	if _, err = g.Write(b); err != nil {
		slog.Error(4, err)
		return
	}
	if err = g.Close(); err != nil {
		slog.Error(5, err)
		return
	}
	req, err := http.NewRequest("POST", tsdbURL, &buf)
	if err != nil {
		slog.Error(6, err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Content-Encoding", "gzip")
	resp, err := client.Do(req)
	if err == nil {
		defer resp.Body.Close()
	}
	// Some problem with connecting to the server; retry later.
	if err != nil || resp.StatusCode != http.StatusNoContent {
		if err != nil {
			slog.Error(7, err)
		} else if resp.StatusCode != http.StatusNoContent {
			slog.Errorln(8, resp.Status)
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				slog.Error(1, err)
			}
			if len(body) > 0 {
				slog.Error(2, string(body))
			}
		}
		t := time.Now().Add(-time.Minute * 30).Unix()
		old := 0
		restored := 0
		for _, dp := range batch {
			if dp.Timestamp < t {
				old++
				continue
			}
			restored++
			tchan <- dp
		}
		if old > 0 {
			slog.Infof("removed %d old records", old)
		}
		d := time.Second * 5
		slog.Infof("restored %d, sleeping %s", restored, d)
		time.Sleep(d)
		return
	} else {
		slog.Infoln("sent", len(batch))
		sent += int64(len(batch))
	}
}
<commit_msg>cmd/scollector: Remove debugs<commit_after>package collect

import (
	"bytes"
	"compress/gzip"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/slog"
)

func send() {
	for {
		qlock.Lock()
		if len(queue) > 0 {
			i := len(queue)
			if i > BatchSize {
				i = BatchSize
			}
			sending := queue[:i]
			queue = queue[i:]
			qlock.Unlock()
			slog.Infof("sending: %d, remaining: %d", len(sending), len(queue))
			sendBatch(sending)
		} else {
			qlock.Unlock()
			time.Sleep(time.Second)
		}
	}
}

func sendBatch(batch opentsdb.MultiDataPoint) {
	b, err := batch.Json()
	if err != nil {
		slog.Error(err)
		// bad JSON encoding, just give up
		return
	}
	var buf bytes.Buffer
	g := gzip.NewWriter(&buf)
	if _, err = g.Write(b); err != nil {
		slog.Error(err)
		return
	}
	if err = g.Close(); err != nil {
		slog.Error(err)
		return
	}
	req, err := http.NewRequest("POST", tsdbURL, &buf)
	if err != nil {
		slog.Error(err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Content-Encoding", "gzip")
	resp, err := client.Do(req)
	if err == nil {
		defer resp.Body.Close()
	}
	// Some problem with connecting to the server; retry later.
	if err != nil || resp.StatusCode != http.StatusNoContent {
		if err != nil {
			slog.Error(err)
		} else if resp.StatusCode != http.StatusNoContent {
			slog.Errorln(resp.Status)
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				slog.Error(err)
			}
			if len(body) > 0 {
				slog.Error(string(body))
			}
		}
		t := time.Now().Add(-time.Minute * 30).Unix()
		old := 0
		restored := 0
		for _, dp := range batch {
			if dp.Timestamp < t {
				old++
				continue
			}
			restored++
			tchan <- dp
		}
		if old > 0 {
			slog.Infof("removed %d old records", old)
		}
		d := time.Second * 5
		slog.Infof("restored %d, sleeping %s", restored, d)
		time.Sleep(d)
		return
	} else {
		slog.Infoln("sent", len(batch))
		sent += int64(len(batch))
	}
}
<|endoftext|>


<commit_before>package bmp_test

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

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	cmds "github.com/cloudfoundry-community/bosh-softlayer-tools/cmds"
	bmp "github.com/cloudfoundry-community/bosh-softlayer-tools/cmds/bmp"
	config "github.com/cloudfoundry-community/bosh-softlayer-tools/config"

	clientsfakes "github.com/cloudfoundry-community/bosh-softlayer-tools/clients/fakes"
)

var _ = Describe("target command", func() {
	var (
		err error

		args    []string
		options cmds.Options
		cmd     cmds.Command

		tmpDir, tmpFileName string

		fakeBmpClient *clientsfakes.FakeBmpClient
	)

	BeforeEach(func() {
		args = []string{"bmp", "target"}
		options = cmds.Options{
			Verbose: false,
			Target:  "http://fake.url",
		}

		tmpDir, err = ioutil.TempDir("", "bmp-target-execute")
		Expect(err).ToNot(HaveOccurred())

		tmpFile, err := ioutil.TempFile(tmpDir, ".bmp_config")
		Expect(err).ToNot(HaveOccurred())

		tmpFileName = tmpFile.Name()

		fakeBmpClient = clientsfakes.NewFakeBmpClient(options.Username, options.Password, options.Target, tmpFileName)
		cmd = bmp.NewTargetCommand(options, fakeBmpClient)
	})

	AfterEach(func() {
		err := os.RemoveAll(tmpDir)
		Expect(err).ToNot(HaveOccurred())
	})

	Describe("NewTargetCommand", func() {
		It("create new TargetCommand", func() {
			Expect(cmd).ToNot(BeNil())

			cmd2 := bmp.NewTargetCommand(options, fakeBmpClient)
			Expect(cmd2).ToNot(BeNil())
			Expect(cmd2).To(Equal(cmd))
		})
	})

	Describe("#Name", func() {
		It("returns the name of a TargetCommand", func() {
			Expect(cmd.Name()).To(Equal("target"))
		})
	})

	Describe("#Description", func() {
		It("returns the description of a TargetCommand", func() {
			Expect(cmd.Description()).To(Equal("Set the URL of Bare Metal Provision Server"))
		})
	})

	Describe("#Usage", func() {
		It("returns the usage text of a TargetCommand", func() {
			Expect(cmd.Usage()).To(Equal("bmp target http://url.to.bmp.server"))
		})
	})

	Describe("#Options", func() {
		It("returns the options of a TargetCommand", func() {
			Expect(cmds.EqualOptions(cmd.Options(), options)).To(BeTrue())
		})
	})

	Describe("#Validate", func() {
		Context("when the options includes a URL", func() {
			It("returns true on validation", func() {
				validate, err := cmd.Validate()
				Expect(validate).To(BeTrue())
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("when the options includes an empty target URL", func() {
			BeforeEach(func() {
				options.Target = ""
				cmd = bmp.NewTargetCommand(options, fakeBmpClient)
			})

			It("returns false on validation", func() {
				validate, err := cmd.Validate()
				Expect(validate).To(BeFalse())
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("when the options includes bad target URL", func() {
			var badCmd1, badCmd2 cmds.Command

			BeforeEach(func() {
				options.Target = "bad-url"
				badCmd1 = bmp.NewTargetCommand(options, fakeBmpClient)

				options.Target = "..."
				badCmd2 = bmp.NewTargetCommand(options, fakeBmpClient)
			})

			It("returns false on validation", func() {
				validate, err := badCmd1.Validate()
				Expect(validate).To(BeFalse())
				Expect(err).To(HaveOccurred())

				validate, err = badCmd2.Validate()
				Expect(validate).To(BeFalse())
				Expect(err).To(HaveOccurred())
			})
		})
	})

	Describe("#Execute", func() {
		It("executes a good TargetCommand", func() {
			rc, err := cmd.Execute(args)

			fmt.Printf("====> err: %#v\n", err)

			Expect(rc).To(Equal(0))
			Expect(err).ToNot(HaveOccurred())

			c := config.NewConfig(fakeBmpClient.ConfigPath())
			configInfo, err := c.LoadConfig()
			Expect(err).ToNot(HaveOccurred())
			Expect(configInfo.TargetUrl).To(Equal("http://fake.url"))
		})
	})
})
<commit_msg>Fixed issue #25 and #26<commit_after>package bmp_test

import (
	"io/ioutil"
	"os"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	cmds "github.com/cloudfoundry-community/bosh-softlayer-tools/cmds"
	bmp "github.com/cloudfoundry-community/bosh-softlayer-tools/cmds/bmp"
	config "github.com/cloudfoundry-community/bosh-softlayer-tools/config"

	clientsfakes "github.com/cloudfoundry-community/bosh-softlayer-tools/clients/fakes"
)

var _ = Describe("target command", func() {
	var (
		err error

		args    []string
		options cmds.Options
		cmd     cmds.Command

		tmpDir, tmpFileName string

		fakeBmpClient *clientsfakes.FakeBmpClient
	)

	BeforeEach(func() {
		args = []string{"bmp", "target"}
		options = cmds.Options{
			Verbose: false,
			Target:  "http://fake.url",
		}

		tmpDir, err = ioutil.TempDir("", "bmp-target-execute")
		Expect(err).ToNot(HaveOccurred())

		tmpFile, err := ioutil.TempFile(tmpDir, ".bmp_config")
		Expect(err).ToNot(HaveOccurred())

		tmpFileName = tmpFile.Name()

		err = ioutil.WriteFile(tmpFileName, []byte("{}"), 0666)
		Expect(err).ToNot(HaveOccurred())

		fakeBmpClient = clientsfakes.NewFakeBmpClient(options.Username, options.Password, options.Target, tmpFileName)
		cmd = bmp.NewTargetCommand(options, fakeBmpClient)
	})

	AfterEach(func() {
		err := os.RemoveAll(tmpDir)
		Expect(err).ToNot(HaveOccurred())
	})

	Describe("NewTargetCommand", func() {
		It("create new TargetCommand", func() {
			Expect(cmd).ToNot(BeNil())

			cmd2 := bmp.NewTargetCommand(options, fakeBmpClient)
			Expect(cmd2).ToNot(BeNil())
			Expect(cmd2).To(Equal(cmd))
		})
	})

	Describe("#Name", func() {
		It("returns the name of a TargetCommand", func() {
			Expect(cmd.Name()).To(Equal("target"))
		})
	})

	Describe("#Description", func() {
		It("returns the description of a TargetCommand", func() {
			Expect(cmd.Description()).To(Equal("Set the URL of Bare Metal Provision Server"))
		})
	})

	Describe("#Usage", func() {
		It("returns the usage text of a TargetCommand", func() {
			Expect(cmd.Usage()).To(Equal("bmp target http://url.to.bmp.server"))
		})
	})

	Describe("#Options", func() {
		It("returns the options of a TargetCommand", func() {
			Expect(cmds.EqualOptions(cmd.Options(), options)).To(BeTrue())
		})
	})

	Describe("#Validate", func() {
		Context("when the options includes a URL", func() {
			It("returns true on validation", func() {
				validate, err := cmd.Validate()
				Expect(validate).To(BeTrue())
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("when the options includes an empty target URL", func() {
			BeforeEach(func() {
				options.Target = ""
				cmd = bmp.NewTargetCommand(options, fakeBmpClient)
			})

			It("returns false on validation", func() {
				validate, err := cmd.Validate()
				Expect(validate).To(BeFalse())
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("when the options includes bad target URL", func() {
			var badCmd1, badCmd2 cmds.Command

			BeforeEach(func() {
				options.Target = "bad-url"
				badCmd1 = bmp.NewTargetCommand(options, fakeBmpClient)

				options.Target = "..."
				badCmd2 = bmp.NewTargetCommand(options, fakeBmpClient)
			})

			It("returns false on validation", func() {
				validate, err := badCmd1.Validate()
				Expect(validate).To(BeFalse())
				Expect(err).To(HaveOccurred())

				validate, err = badCmd2.Validate()
				Expect(validate).To(BeFalse())
				Expect(err).To(HaveOccurred())
			})
		})
	})

	Describe("#Execute", func() {
		It("executes a good TargetCommand", func() {
			rc, err := cmd.Execute(args)
			Expect(rc).To(Equal(0))
			Expect(err).ToNot(HaveOccurred())

			c := config.NewConfig(fakeBmpClient.ConfigPath())
			configInfo, err := c.LoadConfig()
			Expect(err).ToNot(HaveOccurred())
			Expect(configInfo.TargetUrl).To(Equal("http://fake.url"))
		})
	})
})
<|endoftext|>


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

// legacy_boot allows to handover a system running linuxboot/u-root
// to a legacy preinstalled operating system by replacing the traditional
// bootloader path

//
// Synopsis:
//	legacy_boot
//
// Description:
//	If returns to u-root shell, the code didn't found a local bootable option
//
// Notes:
//	The code is looking for boot/grub/grub.cfg file as to identify the
//	boot option.
//	The first bootable device found in the block device tree is the one used
//	Windows is not supported (that is a work in progress)
//
// Example:
//	legacy_boot -v 	- Start the script in verbose mode for debugging purpose

package main

import (
	"bufio"
	"fmt"
	"github.com/u-root/u-root/pkg/kexec"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"flag"
	"syscall"
)

var verbose bool
type options struct {
        verbose bool
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func blkDevicesList(blkpath string, devpath string) []string {
	var blkDevices []string
	files, err := ioutil.ReadDir(blkpath)
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		deviceEntry, err := ioutil.ReadDir(blkpath + file.Name() + devpath)
		if err != nil {
			println("can t read directory")
		}
		blkDevices = append(blkDevices, deviceEntry[0].Name())
	}
	return (blkDevices)
}

// checkForBootableMbr is looking for bootable MBR signature
// Current support is limited to Hard disk devices and USB devices
func checkForBootableMbr(path string) int {
	var b511, b510 byte
	var err error
	f, err := os.Open(path)
	check(err)
	b1 := make([]byte, 512)
	_,err=f.Read(b1)
	check(err)
	b511 = b1[511]
	b510 = b1[510]
	err=f.Close()
	check(err)
	if ((b511 == 0xaa) && (b510 == 0x55)) == true {
		return 1
	}
	return 0
}

// getDevicePartList returns all devices attached to a specific name like /dev/sdaX where X can move from 0 to 127
// FIXME no support for devices which are included into subdirectory within /dev
func getDevicePartList(path string) []string {
	var returnValue []string
	files, err := ioutil.ReadDir("/dev/")
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		if strings.HasPrefix(file.Name(), path) {
			// We shall not return full device name
			if file.Name() != path {
				// We shall check that the remaining part is a number
				returnValue = append(returnValue, file.Name())
			}
		}
	}
	return returnValue
}

// getSupportedFilesystem returns all block file system supported by the linuxboot kernel
func getSupportedFilesystem() []string {
	var returnValue []string
	var err error
	file, err := os.Open("/proc/filesystems")
	check(err)
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		log.Fatal(err)
	}
	var string1, string2 string
	for scanner.Scan() {
		string1 = ""
		string2 = ""
		fmt.Sscanf(scanner.Text(), "%s %s", &string1, &string2)
		if string2 == "" {
			if string1 != "" {
				returnValue = append(returnValue, string1)
			}
		}
	}
	err=file.Close()
	check(err)
	return returnValue

}

// mountEntry tries to mount a specific block device
func mountEntry(path string, supportedFilesystem []string) bool {
	var returnValue bool
	var err error
	err=syscall.Mkdir("/u-root", 0777)
	check(err)
	var flags uintptr
	// Was supposed to be unecessary for kernel 4.x.x
	if verbose {
		println("/dev/" + path)
	}
	for _, filesystem := range supportedFilesystem {
		flags = syscall.MS_MGC_VAL
		// Need to load the filesystem kind supported
		err = syscall.Mkdir("/u-root/"+path, 0777)
		check(err)
		err := syscall.Mount("/dev/"+path, "/u-root/"+path, filesystem, flags, "")
		if err == nil {
			return true
		}
	}
	returnValue = false
	return returnValue
}

func umountEntry(path string) bool {
	var returnValue bool
	var flags int
	// Was supposed to be unecessary for kernel 4.x.x
	flags = syscall.MNT_DETACH
	err := syscall.Unmount(path, flags)
	if err == nil {
		return true
	}
	returnValue = false
	return returnValue
}

// checkBootEntry is looking for grub.cfg file
// and return absolute path to it
func checkBootEntry(mountPoint string) string {
	_, err := os.Stat(mountPoint + "/boot")
	if err == nil {
		// The boot directory is there
		_, err2 := os.Stat(mountPoint + "/boot" + "/grub")
		if err2 == nil {
			_, err3 := os.Stat(mountPoint + "/boot" + "/grub" + "/grub.cfg")
			if err3 == nil {
				// found
				return (mountPoint + "/boot" + "/grub")
			}
		}
	}
	return ""

}

// getFileMenuContent is parsing a grub.cfg file
// input: absolute directory path to grub.cfg
// output: Return a list of strings with the following format
//	 line[3*x] - menuconfig
//	 line[3*x+1] - linux kernel + boot options
// 	 line[3*x+2] - initrd
// and the default boot entry configured into grub.cfg
func getFileMenuContent(path string) ([]string, int) {
	var returnValue []string
	var err error
	file, err := os.Open(path + "/grub.cfg")
	check(err)
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		log.Fatal(err)
	}
	var status int
	var intReturn int
	intReturn = 0
	status = 0
	// When status = 0 we are looking for a menu entry
	// When status = 1 we are looking for a linux entry
	// When status = 2 we are looking for a initrd entry
	var trimmedLine string
	for scanner.Scan() {
		trimmedLine = strings.TrimSpace(scanner.Text())
		trimmedLine = strings.Join(strings.Fields(trimmedLine), " ")
		if !strings.HasPrefix(trimmedLine, "#") {
			if (strings.HasPrefix(trimmedLine, "set default=") ) && (status == 0) {
				fmt.Sscanf(trimmedLine, "set default=\"%d\"", &intReturn)
			}
			if (strings.HasPrefix(trimmedLine, "menuentry ")) && (status == 0) {
				status = 1
				returnValue = append(returnValue, trimmedLine)
			}
			if (strings.HasPrefix(trimmedLine, "linux ")) && (status == 1) {
				status = 2
				returnValue = append(returnValue, trimmedLine)
			}
			if (strings.HasPrefix(trimmedLine, "initrd ")) && (status == 2) {
				status = 0
				returnValue = append(returnValue, trimmedLine)
			}
		}
	}
	err = file.Close()
	check(err)
	return returnValue, intReturn

}

func copyLocal(path string) string {
	var dest string
	var err error
	result := strings.Split(path, "/")
	for _, entry := range result {
		dest = entry
	}
	dest = "/tmp/" + dest
	srcFile, err := os.Open(path)
	check(err)

	destFile, err := os.Create(dest) // creates if file doesn't exist
	check(err)

	_, err = io.Copy(destFile, srcFile) // check first var for number of bytes copied
	check(err)

	err = destFile.Sync()
	check(err)
	err = destFile.Close()
	check(err)
	err = srcFile.Close()
	check(err)
	return dest
}

// kexecEntry is booting new kernel based on the content of grub.cfg
func kexecEntry(grubConfPath string, mountPoint string) {
	var fileMenuContent []string
	var entry int
	var localKernelPath string
	var localInitrdPath string
	if verbose {
		println(grubConfPath)
	}
	fileMenuContent, entry = getFileMenuContent(grubConfPath)
	var kernel string
	var kernelParameter string
	var initrd string
	var kernelInfos []string
	kernelInfos = strings.Fields(fileMenuContent[3*entry+1])
	kernel = kernelInfos[1]
	var count int
	count = 0
	for _, field := range kernelInfos {
		if count > 1 {
			kernelParameter = kernelParameter + " " + field
		}
		count = count + 1
	}
	fmt.Sscanf(fileMenuContent[3*entry+2], "initrd %s", &initrd)
	if verbose {
		println("************** boot parameters  ********************")
		println(kernel)
		println(kernelParameter)
		println(initrd)
		println("****************************************************")
	}
	localKernelPath = copyLocal(mountPoint + kernel)
	localInitrdPath = copyLocal(mountPoint + initrd)
	if verbose {
		println(localKernelPath)
	}
	umountEntry(mountPoint)
	// We can kexec the kernel with localKernelPath as kernel entry, kernelParameter as parameter and initrd as initrd !
	log.Printf("Loading %s for kernel\n", localKernelPath)

	kernelDesc, err := os.OpenFile(localKernelPath, os.O_RDONLY, 0)
	if err != nil {
		log.Fatalf("open(%q): %v", localKernelPath, err)
	}
	// defer kernelDesc.Close()

	var ramfs *os.File
	ramfs, err = os.OpenFile(localInitrdPath, os.O_RDONLY, 0)
	if err != nil {
		log.Fatalf("open(%q): %v", localInitrdPath, err)
	}
	// defer ramfs.Close()

	if err := kexec.FileLoad(kernelDesc, ramfs, kernelParameter); err != nil {
		log.Fatalf("%v", err)
	}
 	err = ramfs.Close()
	check(err)
	err = kernelDesc.Close()
	check(err)
	if err := kexec.Reboot(); err != nil {
		log.Fatalf("%v", err)
	}

}

func registerFlags(f *flag.FlagSet) *options {
        o := &options{}
        f.BoolVar(&o.verbose, "v", false, "Set verbose output")
        return o
}

func main() {
	var blkList []string
	var supportedFilesystem []string
	verbose = false
        opts := registerFlags(flag.CommandLine)
	flag.Parse()

        if opts.verbose != false {
		verbose = true
        }

	supportedFilesystem = getSupportedFilesystem()
	if verbose {
		println("************** Supported Filesystem by current linuxboot ********************")
		for _, filesystem := range supportedFilesystem {
			println(filesystem)
		}
		println("*****************************************************************************")
	}
	blkList = blkDevicesList("/sys/dev/block/", "/device/block/")
	// We must validate if the MBR is bootable or not and keep the
	// devices which do have such support
	// drive are easy to detect
	for _, entry := range blkList {
		if checkForBootableMbr("/dev/"+entry) == 1 {
			fmt.Println("Bootable device found")
			// We need to loop on the device entries which are into /dev/<device>X
			// and mount each partitions as to find /boot entry if it is available somewhere
			var devicePartList []string
			devicePartList = getDevicePartList(entry)
			for _, deviceList := range devicePartList {
				if mountEntry(deviceList, supportedFilesystem) {
					if verbose {
						println("mount succeed")
					}
					var grubConfPath = checkBootEntry("/u-root/" + deviceList)
					if grubConfPath != "" {
						if verbose {
							println("calling basic kexec")
						}
						kexecEntry(grubConfPath, "/u-root/"+deviceList)
					}
				}
				umountEntry("/u-root/" + deviceList)
			}
		}
	}
	println("Sorry no bootable device found")
}
<commit_msg>use Fields into getSupportedFilesystem instead of string parsing<commit_after>// Copyright 2012-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// legacy_boot allows to handover a system running linuxboot/u-root
// to a legacy preinstalled operating system by replacing the traditional
// bootloader path

//
// Synopsis:
//	legacy_boot
//
// Description:
//	If returns to u-root shell, the code didn't found a local bootable option
//
// Notes:
//	The code is looking for boot/grub/grub.cfg file as to identify the
//	boot option.
//	The first bootable device found in the block device tree is the one used
//	Windows is not supported (that is a work in progress)
//
// Example:
//	legacy_boot -v 	- Start the script in verbose mode for debugging purpose

package main

import (
	"bufio"
	"fmt"
	"github.com/u-root/u-root/pkg/kexec"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"flag"
	"syscall"
)

var verbose bool
type options struct {
        verbose bool
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func blkDevicesList(blkpath string, devpath string) []string {
	var blkDevices []string
	files, err := ioutil.ReadDir(blkpath)
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		deviceEntry, err := ioutil.ReadDir(blkpath + file.Name() + devpath)
		if err != nil {
			println("can t read directory")
		}
		blkDevices = append(blkDevices, deviceEntry[0].Name())
	}
	return (blkDevices)
}

// checkForBootableMbr is looking for bootable MBR signature
// Current support is limited to Hard disk devices and USB devices
func checkForBootableMbr(path string) int {
	var b511, b510 byte
	var err error
	f, err := os.Open(path)
	check(err)
	b1 := make([]byte, 512)
	_,err=f.Read(b1)
	check(err)
	b511 = b1[511]
	b510 = b1[510]
	err=f.Close()
	check(err)
	if ((b511 == 0xaa) && (b510 == 0x55)) == true {
		return 1
	}
	return 0
}

// getDevicePartList returns all devices attached to a specific name like /dev/sdaX where X can move from 0 to 127
// FIXME no support for devices which are included into subdirectory within /dev
func getDevicePartList(path string) []string {
	var returnValue []string
	files, err := ioutil.ReadDir("/dev/")
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		if strings.HasPrefix(file.Name(), path) {
			// We shall not return full device name
			if file.Name() != path {
				// We shall check that the remaining part is a number
				returnValue = append(returnValue, file.Name())
			}
		}
	}
	return returnValue
}

// getSupportedFilesystem returns all block file system supported by the linuxboot kernel
func getSupportedFilesystem() []string {
	var returnValue []string
	var err error
	file, err := os.Open("/proc/filesystems")
	check(err)
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		log.Fatal(err)
	}

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if fields[0] == "nodev" {
			continue
		}
		if fields[1] != "" {
			returnValue = append(returnValue, fields[1])
    		}
	}

	err=file.Close()
	check(err)
	return returnValue

}

// mountEntry tries to mount a specific block device
func mountEntry(path string, supportedFilesystem []string) bool {
	var returnValue bool
	var err error
	err=syscall.Mkdir("/u-root", 0777)
	check(err)
	var flags uintptr
	// Was supposed to be unecessary for kernel 4.x.x
	if verbose {
		println("/dev/" + path)
	}
	for _, filesystem := range supportedFilesystem {
		flags = syscall.MS_MGC_VAL
		// Need to load the filesystem kind supported
		err = syscall.Mkdir("/u-root/"+path, 0777)
		check(err)
		err := syscall.Mount("/dev/"+path, "/u-root/"+path, filesystem, flags, "")
		if err == nil {
			return true
		}
	}
	returnValue = false
	return returnValue
}

func umountEntry(path string) bool {
	var returnValue bool
	var flags int
	// Was supposed to be unecessary for kernel 4.x.x
	flags = syscall.MNT_DETACH
	err := syscall.Unmount(path, flags)
	if err == nil {
		return true
	}
	returnValue = false
	return returnValue
}

// checkBootEntry is looking for grub.cfg file
// and return absolute path to it
func checkBootEntry(mountPoint string) string {
	_, err := os.Stat(mountPoint + "/boot")
	if err == nil {
		// The boot directory is there
		_, err2 := os.Stat(mountPoint + "/boot" + "/grub")
		if err2 == nil {
			_, err3 := os.Stat(mountPoint + "/boot" + "/grub" + "/grub.cfg")
			if err3 == nil {
				// found
				return (mountPoint + "/boot" + "/grub")
			}
		}
	}
	return ""

}

// getFileMenuContent is parsing a grub.cfg file
// input: absolute directory path to grub.cfg
// output: Return a list of strings with the following format
//	 line[3*x] - menuconfig
//	 line[3*x+1] - linux kernel + boot options
// 	 line[3*x+2] - initrd
// and the default boot entry configured into grub.cfg
func getFileMenuContent(path string) ([]string, int) {
	var returnValue []string
	var err error
	file, err := os.Open(path + "/grub.cfg")
	check(err)
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		log.Fatal(err)
	}
	var status int
	var intReturn int
	intReturn = 0
	status = 0
	// When status = 0 we are looking for a menu entry
	// When status = 1 we are looking for a linux entry
	// When status = 2 we are looking for a initrd entry
	var trimmedLine string
	for scanner.Scan() {
		trimmedLine = strings.TrimSpace(scanner.Text())
		trimmedLine = strings.Join(strings.Fields(trimmedLine), " ")
		if !strings.HasPrefix(trimmedLine, "#") {
			if (strings.HasPrefix(trimmedLine, "set default=") ) && (status == 0) {
				fmt.Sscanf(trimmedLine, "set default=\"%d\"", &intReturn)
			}
			if (strings.HasPrefix(trimmedLine, "menuentry ")) && (status == 0) {
				status = 1
				returnValue = append(returnValue, trimmedLine)
			}
			if (strings.HasPrefix(trimmedLine, "linux ")) && (status == 1) {
				status = 2
				returnValue = append(returnValue, trimmedLine)
			}
			if (strings.HasPrefix(trimmedLine, "initrd ")) && (status == 2) {
				status = 0
				returnValue = append(returnValue, trimmedLine)
			}
		}
	}
	err = file.Close()
	check(err)
	return returnValue, intReturn

}

func copyLocal(path string) string {
	var dest string
	var err error
	result := strings.Split(path, "/")
	for _, entry := range result {
		dest = entry
	}
	dest = "/tmp/" + dest
	srcFile, err := os.Open(path)
	check(err)

	destFile, err := os.Create(dest) // creates if file doesn't exist
	check(err)

	_, err = io.Copy(destFile, srcFile) // check first var for number of bytes copied
	check(err)

	err = destFile.Sync()
	check(err)
	err = destFile.Close()
	check(err)
	err = srcFile.Close()
	check(err)
	return dest
}

// kexecEntry is booting new kernel based on the content of grub.cfg
func kexecEntry(grubConfPath string, mountPoint string) {
	var fileMenuContent []string
	var entry int
	var localKernelPath string
	var localInitrdPath string
	if verbose {
		println(grubConfPath)
	}
	fileMenuContent, entry = getFileMenuContent(grubConfPath)
	var kernel string
	var kernelParameter string
	var initrd string
	var kernelInfos []string
	kernelInfos = strings.Fields(fileMenuContent[3*entry+1])
	kernel = kernelInfos[1]
	var count int
	count = 0
	for _, field := range kernelInfos {
		if count > 1 {
			kernelParameter = kernelParameter + " " + field
		}
		count = count + 1
	}
	fmt.Sscanf(fileMenuContent[3*entry+2], "initrd %s", &initrd)
	if verbose {
		println("************** boot parameters  ********************")
		println(kernel)
		println(kernelParameter)
		println(initrd)
		println("****************************************************")
	}
	localKernelPath = copyLocal(mountPoint + kernel)
	localInitrdPath = copyLocal(mountPoint + initrd)
	if verbose {
		println(localKernelPath)
	}
	umountEntry(mountPoint)
	// We can kexec the kernel with localKernelPath as kernel entry, kernelParameter as parameter and initrd as initrd !
	log.Printf("Loading %s for kernel\n", localKernelPath)

	kernelDesc, err := os.OpenFile(localKernelPath, os.O_RDONLY, 0)
	if err != nil {
		log.Fatalf("open(%q): %v", localKernelPath, err)
	}
	// defer kernelDesc.Close()

	var ramfs *os.File
	ramfs, err = os.OpenFile(localInitrdPath, os.O_RDONLY, 0)
	if err != nil {
		log.Fatalf("open(%q): %v", localInitrdPath, err)
	}
	// defer ramfs.Close()

	if err := kexec.FileLoad(kernelDesc, ramfs, kernelParameter); err != nil {
		log.Fatalf("%v", err)
	}
 	err = ramfs.Close()
	check(err)
	err = kernelDesc.Close()
	check(err)
	if err := kexec.Reboot(); err != nil {
		log.Fatalf("%v", err)
	}

}

func registerFlags(f *flag.FlagSet) *options {
        o := &options{}
        f.BoolVar(&o.verbose, "v", false, "Set verbose output")
        return o
}

func main() {
	var blkList []string
	var supportedFilesystem []string
	verbose = false
        opts := registerFlags(flag.CommandLine)
	flag.Parse()

        if opts.verbose != false {
		verbose = true
        }

	supportedFilesystem = getSupportedFilesystem()
	if verbose {
		println("************** Supported Filesystem by current linuxboot ********************")
		for _, filesystem := range supportedFilesystem {
			println(filesystem)
		}
		println("*****************************************************************************")
	}
	blkList = blkDevicesList("/sys/dev/block/", "/device/block/")
	// We must validate if the MBR is bootable or not and keep the
	// devices which do have such support
	// drive are easy to detect
	for _, entry := range blkList {
		if checkForBootableMbr("/dev/"+entry) == 1 {
			fmt.Println("Bootable device found")
			// We need to loop on the device entries which are into /dev/<device>X
			// and mount each partitions as to find /boot entry if it is available somewhere
			var devicePartList []string
			devicePartList = getDevicePartList(entry)
			for _, deviceList := range devicePartList {
				if mountEntry(deviceList, supportedFilesystem) {
					if verbose {
						println("mount succeed")
					}
					var grubConfPath = checkBootEntry("/u-root/" + deviceList)
					if grubConfPath != "" {
						if verbose {
							println("calling basic kexec")
						}
						kexecEntry(grubConfPath, "/u-root/"+deviceList)
					}
				}
				umountEntry("/u-root/" + deviceList)
			}
		}
	}
	println("Sorry no bootable device found")
}
<|endoftext|>


<commit_before>// etcd will queried when this store is used

package etcdstore

import (
	"errors"
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"log"
	"strings"
)

const _leaf = "_leaf"
const _sep = "\t"
const _roptimize = "/_roptimize"

type EtcdStore struct {
	hosts      []string     // http://host1:port,..
	ROptimize  bool         // reverse lookup optimization
	FastLookup bool         // fast return, will return the first match
	client     *etcd.Client // etcd connection object
	storenode  string       // path to where the range store is etcd
}

// Connect to the Etcd Store
// 1. make sure Etcd Store is running
// 2. check whether _range_store value is set to 'loaded'
// 3. create the interface for client connections
func ConnectEtcdStore(hosts []string, roptimize, fast bool, node string) (e *EtcdStore, err error) {
	if len(node) > 0 && node[len(node)-1] == '/' {
		node = node[:len(node)-1]
	}
	client := etcd.NewClient(hosts)
	e = &EtcdStore{hosts: hosts, ROptimize: roptimize, FastLookup: fast, client: client, storenode: node}
	// test the consistency of etcd store
	var obj = fmt.Sprintf("%s/%s", e.storenode, "_range_store")
	var response *etcd.Response
	response, err = client.Get(obj, false, false)
	// we have to make sure we have a clean etcd store
	if err != nil && err.(*etcd.EtcdError).ErrorCode == 100 {
		log.Printf("ERROR: Looks like the etcd-store is not loaded with data to serve as rangestore [Key: %s, Value: NOT FOUND]\n", obj)
		return nil, errors.New("ERROR: etcd-store is NOT LOADED with data")
	} else if err != nil {
		log.Println("ERROR: Etcd Store Returned Error")
		log.Println(err)
		return nil, err
	} else if err == nil {
		// if we get no error, it means key if present. We need to make sure it is set to 'loaded'
		if response.Node.Value != "loaded" {
			log.Printf("ERROR: Looks like the etcd-store is not ready to serve as rangestore [Key: %s, Value: %s]\n", response.Node.Key, response.Node.Value)
			return nil, errors.New("ERROR: etcd-store is NOT READY to serve")
		}
	}

	return e, nil
}

func (e *EtcdStore) DisconnectEtcdStore() {
	e.client.Close()
	return
}

////////////////////
// LOOKUP CLUSTER //
////////////////////

// LOGIC
// -----
// * for the first element in cluster create results array
//   * check whether the cluster is a leaf node
//   * if yes, call KeyLookup, with key == NODES
//   * if not, call listClusters
// * if more elements are there, repeat the above
//   (unlike, filestore we don't call ArrayToSet since etcd
//   is populated by a program, not by morals)
func (e *EtcdStore) ClusterLookup(cluster *[]string) (*[]string, error) {
	// store the resuls
	var results = make([]string, 0)
	// for each cluster, do a lookup
	// (this will only happen only for nested lookups eg, %%..)
	for _, elem := range *cluster {
		// handle RANGE separately
		if elem == "RANGE" {
			elem = "/"
		}
		var err error
		isLeaf, err := e.checkIsLeafNode(elem)
		if err != nil {
			return &[]string{}, err
		}
		// if it is a leaf node, we need do a KeyLookup (NODES)
		if isLeaf {
			// by default, lookup for NODES
			result, err := e.KeyLookup(&[]string{elem}, "NODES")
			if err != nil {
				return &[]string{}, err
			}
			results = append(results, *result...)
		} else { // we need to return the children
			result, err := e.listClusters(elem)
			if err != nil {
				return &[]string{}, err
			}
			results = append(results, result...)
		}

	}

	return &results, nil
}

func (e *EtcdStore) KeyLookup(cluster *[]string, key string) (*[]string, error) {
	// store the resuls
	var results = make([]string, 0)
	// this will most likely be single element arrays
	// can't think of a reason otherwise
	for _, elem := range *cluster {
		dir := e.clusterToPath(elem)
		node := fmt.Sprintf("%s/%s", dir, key)
		var result []string
		if key == "KEYS" {
			response, _, _, found, err := e.retrieveFromEtcd(dir, false, false)
			// if there is an error, return err
			if err != nil { // got error
				return &[]string{}, err
			} else if !found { // key not found
				return &[]string{}, errors.New(fmt.Sprintf("KeyLookup for [%s:%s] Failed (Error: No KEY Found)", elem, key))
			}

			// if response is NOT for a dir
			if !response.Node.Dir {
				var _err = fmt.Sprintf("Expected value of lookup [%s] to be a dir, recieved a leaf file", dir)
				log.Printf(_err)
				return &[]string{}, errors.New(_err)
			}

			for _, n := range response.Node.Nodes {
				// replace '/' with '-', also root will always be '/'
				_node := strings.Split(n.Key, "/")
				if len(_node) > 0 {
					result = append(result, _node[len(_node)-1])
				}
			}
		} else {
			// 1. read the key in etcd
			// 2. append the result
			_, _, value, found, err := e.retrieveFromEtcd(node, false, false)
			if err != nil {
				return &[]string{}, errors.New(fmt.Sprintf("KeyLookup for [%s:%s] Failed (Error: %s)", elem, key, err))
			} else if !found {
				return &[]string{}, errors.New(fmt.Sprintf("KeyLookup for [%s:%s] Failed (Error: No KEY Found)", elem, key))
			}
			result = strings.Split(value, _sep)
		}
		// append the result with results
		results = append(results, result...)
	}

	return &results, nil
}

////////////////////
// LOOKUP REVERSE //
////////////////////

// same as KeyReverseLookupAttr where attr == NODES
func (e *EtcdStore) KeyReverseLookup(key string) (*[]string, error) {
	return e.KeyReverseLookupAttr(key, "NODES")
}

// same as KeyReverseLookupAttr where attr == NODES and hint == ""
func (e *EtcdStore) KeyReverseLookupAttr(key string, attr string) (*[]string, error) {
	// optimization, for nodes don't do the tough thing
	if e.ROptimize && attr == "NODES" {
		return e.optimizedNodeReverseLookup(key)
	}
	return e.KeyReverseLookupHint(key, attr, "")
}

// given a key, it will search for the cluster where the attr has that key,
// hint is to limit the scope of search
func (e *EtcdStore) KeyReverseLookupHint(key string, attr string, hint string) (*[]string, error) {
	var clusters *[]string
	var err error
	var results = make([]string, 0)
	var seen bool

	clusters, err = e.getAllLeafNodes(hint)
	if err != nil {
		return &results, nil
	}

	for _, elem := range *clusters {

		// 1. read the key in etcd
		// 2. append the result
		_, _, value, found, err := e.retrieveFromEtcd(fmt.Sprintf("%s/%s", e.clusterToPath(elem), attr), false, false)
		if err != nil {
			return &[]string{}, errors.New(fmt.Sprintf("KeyLookup for [%s:%s] Failed (Error: %s)", elem, key, err))
		} else if !found {
			continue
		} else {
			result := strings.Split(value, _sep)
			for _, i := range result {
				if i == key {
					results = append(results, elem)
					seen = true
					break
				}
			}
			if seen && e.FastLookup {
				return &results, nil
			}
		}
	}
	return &results, nil
}

////////////////////////
// Internal Functions //
////////////////////////

// given a cluster name, it will convert to cluster
// in the file system
func (e *EtcdStore) clusterToPath(cluster string) string {
	return fmt.Sprintf("%s/%s", e.storenode, strings.Replace(cluster, "-", "/", -1))
}

// Get all the leaf cluster nodes for a given dir
// it is not efficient since we have to walk down all the path
func (e *EtcdStore) getAllLeafNodes(root string) (*[]string, error) {
	var results = make([]string, 0)
	var err error

	// root a leaf node
	isleaf, err := e.checkIsLeafNode(root)
	if err != nil {
		return &[]string{}, err
	}
	if isleaf {
		return &[]string{root}, nil
	}

	err = e._getAllLeafNodes(e.clusterToPath(root), &results)

	if err != nil {
		return &[]string{}, err
	}

	// fix the results in place
	for i := 0; i < len(results); i++ {
		results[i] = strings.Replace(strings.Trim(results[i], "/"), "/", "-", -1)
	}

	return &results, nil
}

// the function that really does the work
func (e *EtcdStore) _getAllLeafNodes(root string, results *[]string) error {
	response, _, _, found, err := e.retrieveFromEtcd(root, false, false)
	// if there is an error, return err
	if err != nil { // got error
		return err
	} else if !found { // key not found
		return errors.New(fmt.Sprintf("DirLookup for [%s] Failed (Error: No DIR Found)", root))
	}

	// if response is NOT for a dir
	if !response.Node.Dir {
		var _err = fmt.Sprintf("Expected value of lookup [%s] to be a dir, recieved a leaf file", root)
		log.Printf(_err)
		return errors.New(_err)
	}

	for _, n := range response.Node.Nodes {
		status, err := e.checkIsLeafNode(n.Key)
		if err != nil {
			return err
		}
		if status {
			*results = append(*results, n.Key)
		} else {
			_ = e._getAllLeafNodes(n.Key, results)
		}
	}

	return nil
}

// reads the child clusters of this cluster.
// returns only those nodes for which this cluster is parent
func (e *EtcdStore) listClusters(cluster string) ([]string, error) {
	var dir = e.clusterToPath(cluster)
	var children = make([]string, 0)
	// list the nodes under this cluster.
	// we are sure when this call was made, the check
	// has been made to sure this is not a leaf node
	response, _, _, found, err := e.retrieveFromEtcd(dir, false, false)
	// if there is an error, return err
	if err != nil { // got error
		return []string{}, err
	} else if !found { // key not found
		return []string{}, nil
	}

	// if response is NOT for a dir
	if !response.Node.Dir {
		var _err = fmt.Sprintf("Expected value of lookup [%s] to be a dir, recieved a leaf file", dir)
		log.Printf(_err)
		return []string{}, errors.New(_err)
	}

	for _, n := range response.Node.Nodes {
		// replace '/' with '-', also root will always be '/'
		_node := strings.Replace(n.Key[1:], "/", "-", -1)
		children = append(children, _node)
	}

	return children, nil
}

// Checks whether the cluster is in leaf or not
// It will return error if the cluster doesn't exist,
// false if not a leaf node, true otherwise
func (e *EtcdStore) checkIsLeafNode(cluster string) (found bool, err error) {
	var dir = e.clusterToPath(cluster)
	_, _, _, found, err = e.retrieveFromEtcd(fmt.Sprintf("%s/%s", dir, _leaf), false, false)
	// found and err are set properly by retrieveFromEtcd
	return found, err
}

// given a object, return the response from the etcd cluster
func (e *EtcdStore) retrieveFromEtcd(object string, sort, recursive bool) (response *etcd.Response, key string, value string, found bool, err error) {
	response, err = e.client.Get(object, sort, recursive)

	// Check whether the error is Key NOT Found
	if err != nil && err.(*etcd.EtcdError).ErrorCode == 100 {
		return response, object, "", false, nil
	} else if err != nil { // error could be cluster is unreachble, etc..
		log.Printf("ERROR: Etcd Store Returned Error: %s\n", err)
		return response, object, "", false, err
	}

	// if all is good
	return response, response.Node.Key, response.Node.Value, true, nil
}

// same as KeyReverseLookupAttr where attr == NODES and hint == ""
func (e *EtcdStore) optimizedNodeReverseLookup(key string) (*[]string, error) {
	_, _, value, found, err := e.retrieveFromEtcd(fmt.Sprintf("%s/%s", _roptimize, key), false, false)
	if !found {
		return &[]string{}, nil
	} else if err != nil {
		return &[]string{}, err
	}
	values := strings.Split(value, _sep)
	return &values, nil
}
<commit_msg>reverse lookup on failure should throw out error<commit_after>// etcd will queried when this store is used

package etcdstore

import (
	"errors"
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"log"
	"strings"
)

const _leaf = "_leaf"
const _sep = "\t"
const _roptimize = "/_roptimize"

type EtcdStore struct {
	hosts      []string     // http://host1:port,..
	ROptimize  bool         // reverse lookup optimization
	FastLookup bool         // fast return, will return the first match
	client     *etcd.Client // etcd connection object
	storenode  string       // path to where the range store is etcd
}

// Connect to the Etcd Store
// 1. make sure Etcd Store is running
// 2. check whether _range_store value is set to 'loaded'
// 3. create the interface for client connections
func ConnectEtcdStore(hosts []string, roptimize, fast bool, node string) (e *EtcdStore, err error) {
	if len(node) > 0 && node[len(node)-1] == '/' {
		node = node[:len(node)-1]
	}
	client := etcd.NewClient(hosts)
	e = &EtcdStore{hosts: hosts, ROptimize: roptimize, FastLookup: fast, client: client, storenode: node}
	// test the consistency of etcd store
	var obj = fmt.Sprintf("%s/%s", e.storenode, "_range_store")
	var response *etcd.Response
	response, err = client.Get(obj, false, false)
	// we have to make sure we have a clean etcd store
	if err != nil && err.(*etcd.EtcdError).ErrorCode == 100 {
		log.Printf("ERROR: Looks like the etcd-store is not loaded with data to serve as rangestore [Key: %s, Value: NOT FOUND]\n", obj)
		return nil, errors.New("ERROR: etcd-store is NOT LOADED with data")
	} else if err != nil {
		log.Println("ERROR: Etcd Store Returned Error")
		log.Println(err)
		return nil, err
	} else if err == nil {
		// if we get no error, it means key if present. We need to make sure it is set to 'loaded'
		if response.Node.Value != "loaded" {
			log.Printf("ERROR: Looks like the etcd-store is not ready to serve as rangestore [Key: %s, Value: %s]\n", response.Node.Key, response.Node.Value)
			return nil, errors.New("ERROR: etcd-store is NOT READY to serve")
		}
	}

	return e, nil
}

func (e *EtcdStore) DisconnectEtcdStore() {
	e.client.Close()
	return
}

////////////////////
// LOOKUP CLUSTER //
////////////////////

// LOGIC
// -----
// * for the first element in cluster create results array
//   * check whether the cluster is a leaf node
//   * if yes, call KeyLookup, with key == NODES
//   * if not, call listClusters
// * if more elements are there, repeat the above
//   (unlike, filestore we don't call ArrayToSet since etcd
//   is populated by a program, not by morals)
func (e *EtcdStore) ClusterLookup(cluster *[]string) (*[]string, error) {
	// store the resuls
	var results = make([]string, 0)
	// for each cluster, do a lookup
	// (this will only happen only for nested lookups eg, %%..)
	for _, elem := range *cluster {
		// handle RANGE separately
		if elem == "RANGE" {
			elem = "/"
		}
		var err error
		isLeaf, err := e.checkIsLeafNode(elem)
		if err != nil {
			return &[]string{}, err
		}
		// if it is a leaf node, we need do a KeyLookup (NODES)
		if isLeaf {
			// by default, lookup for NODES
			result, err := e.KeyLookup(&[]string{elem}, "NODES")
			if err != nil {
				return &[]string{}, err
			}
			results = append(results, *result...)
		} else { // we need to return the children
			result, err := e.listClusters(elem)
			if err != nil {
				return &[]string{}, err
			}
			results = append(results, result...)
		}

	}

	return &results, nil
}

func (e *EtcdStore) KeyLookup(cluster *[]string, key string) (*[]string, error) {
	// store the resuls
	var results = make([]string, 0)
	// this will most likely be single element arrays
	// can't think of a reason otherwise
	for _, elem := range *cluster {
		dir := e.clusterToPath(elem)
		node := fmt.Sprintf("%s/%s", dir, key)
		var result []string
		if key == "KEYS" {
			response, _, _, found, err := e.retrieveFromEtcd(dir, false, false)
			// if there is an error, return err
			if err != nil { // got error
				return &[]string{}, err
			} else if !found { // key not found
				return &[]string{}, errors.New(fmt.Sprintf("KeyLookup for [%s:%s] Failed (Error: No KEY Found)", elem, key))
			}

			// if response is NOT for a dir
			if !response.Node.Dir {
				var _err = fmt.Sprintf("Expected value of lookup [%s] to be a dir, recieved a leaf file", dir)
				log.Printf(_err)
				return &[]string{}, errors.New(_err)
			}

			for _, n := range response.Node.Nodes {
				// replace '/' with '-', also root will always be '/'
				_node := strings.Split(n.Key, "/")
				if len(_node) > 0 {
					result = append(result, _node[len(_node)-1])
				}
			}
		} else {
			// 1. read the key in etcd
			// 2. append the result
			_, _, value, found, err := e.retrieveFromEtcd(node, false, false)
			if err != nil {
				return &[]string{}, errors.New(fmt.Sprintf("KeyLookup for [%s:%s] Failed (Error: %s)", elem, key, err))
			} else if !found {
				return &[]string{}, errors.New(fmt.Sprintf("KeyLookup for [%s:%s] Failed (Error: No KEY Found)", elem, key))
			}
			result = strings.Split(value, _sep)
		}
		// append the result with results
		results = append(results, result...)
	}

	return &results, nil
}

////////////////////
// LOOKUP REVERSE //
////////////////////

// same as KeyReverseLookupAttr where attr == NODES
func (e *EtcdStore) KeyReverseLookup(key string) (*[]string, error) {
	return e.KeyReverseLookupAttr(key, "NODES")
}

// same as KeyReverseLookupAttr where attr == NODES and hint == ""
func (e *EtcdStore) KeyReverseLookupAttr(key string, attr string) (*[]string, error) {
	// optimization, for nodes don't do the tough thing
	if e.ROptimize && attr == "NODES" {
		return e.optimizedNodeReverseLookup(key)
	}
	return e.KeyReverseLookupHint(key, attr, "")
}

// given a key, it will search for the cluster where the attr has that key,
// hint is to limit the scope of search
func (e *EtcdStore) KeyReverseLookupHint(key string, attr string, hint string) (*[]string, error) {
	var clusters *[]string
	var err error
	var results = make([]string, 0)
	var seen bool

	clusters, err = e.getAllLeafNodes(hint)
	if err != nil {
		return &results, nil
	}

	for _, elem := range *clusters {

		// 1. read the key in etcd
		// 2. append the result
		_, _, value, found, err := e.retrieveFromEtcd(fmt.Sprintf("%s/%s", e.clusterToPath(elem), attr), false, false)
		if err != nil {
			return &[]string{}, errors.New(fmt.Sprintf("KeyLookup for [%s:%s] Failed (Error: %s)", elem, key, err))
		} else if !found {
			continue
		} else {
			result := strings.Split(value, _sep)
			for _, i := range result {
				if i == key {
					results = append(results, elem)
					seen = true
					break
				}
			}
			if seen && e.FastLookup {
				return &results, nil
			}
		}
	}
	return &results, nil
}

////////////////////////
// Internal Functions //
////////////////////////

// given a cluster name, it will convert to cluster
// in the file system
func (e *EtcdStore) clusterToPath(cluster string) string {
	return fmt.Sprintf("%s/%s", e.storenode, strings.Replace(cluster, "-", "/", -1))
}

// Get all the leaf cluster nodes for a given dir
// it is not efficient since we have to walk down all the path
func (e *EtcdStore) getAllLeafNodes(root string) (*[]string, error) {
	var results = make([]string, 0)
	var err error

	// root a leaf node
	isleaf, err := e.checkIsLeafNode(root)
	if err != nil {
		return &[]string{}, err
	}
	if isleaf {
		return &[]string{root}, nil
	}

	err = e._getAllLeafNodes(e.clusterToPath(root), &results)

	if err != nil {
		return &[]string{}, err
	}

	// fix the results in place
	for i := 0; i < len(results); i++ {
		results[i] = strings.Replace(strings.Trim(results[i], "/"), "/", "-", -1)
	}

	return &results, nil
}

// the function that really does the work
func (e *EtcdStore) _getAllLeafNodes(root string, results *[]string) error {
	response, _, _, found, err := e.retrieveFromEtcd(root, false, false)
	// if there is an error, return err
	if err != nil { // got error
		return err
	} else if !found { // key not found
		return errors.New(fmt.Sprintf("DirLookup for [%s] Failed (Error: No DIR Found)", root))
	}

	// if response is NOT for a dir
	if !response.Node.Dir {
		var _err = fmt.Sprintf("Expected value of lookup [%s] to be a dir, recieved a leaf file", root)
		log.Printf(_err)
		return errors.New(_err)
	}

	for _, n := range response.Node.Nodes {
		status, err := e.checkIsLeafNode(n.Key)
		if err != nil {
			return err
		}
		if status {
			*results = append(*results, n.Key)
		} else {
			_ = e._getAllLeafNodes(n.Key, results)
		}
	}

	return nil
}

// reads the child clusters of this cluster.
// returns only those nodes for which this cluster is parent
func (e *EtcdStore) listClusters(cluster string) ([]string, error) {
	var dir = e.clusterToPath(cluster)
	var children = make([]string, 0)
	// list the nodes under this cluster.
	// we are sure when this call was made, the check
	// has been made to sure this is not a leaf node
	response, _, _, found, err := e.retrieveFromEtcd(dir, false, false)
	// if there is an error, return err
	if err != nil { // got error
		return []string{}, err
	} else if !found { // key not found
		return []string{}, nil
	}

	// if response is NOT for a dir
	if !response.Node.Dir {
		var _err = fmt.Sprintf("Expected value of lookup [%s] to be a dir, recieved a leaf file", dir)
		log.Printf(_err)
		return []string{}, errors.New(_err)
	}

	for _, n := range response.Node.Nodes {
		// replace '/' with '-', also root will always be '/'
		_node := strings.Replace(n.Key[1:], "/", "-", -1)
		children = append(children, _node)
	}

	return children, nil
}

// Checks whether the cluster is in leaf or not
// It will return error if the cluster doesn't exist,
// false if not a leaf node, true otherwise
func (e *EtcdStore) checkIsLeafNode(cluster string) (found bool, err error) {
	var dir = e.clusterToPath(cluster)
	_, _, _, found, err = e.retrieveFromEtcd(fmt.Sprintf("%s/%s", dir, _leaf), false, false)
	// found and err are set properly by retrieveFromEtcd
	return found, err
}

// given a object, return the response from the etcd cluster
func (e *EtcdStore) retrieveFromEtcd(object string, sort, recursive bool) (response *etcd.Response, key string, value string, found bool, err error) {
	response, err = e.client.Get(object, sort, recursive)

	// Check whether the error is Key NOT Found
	if err != nil && err.(*etcd.EtcdError).ErrorCode == 100 {
		return response, object, "", false, nil
	} else if err != nil { // error could be cluster is unreachble, etc..
		log.Printf("ERROR: Etcd Store Returned Error: %s\n", err)
		return response, object, "", false, err
	}

	// if all is good
	return response, response.Node.Key, response.Node.Value, true, nil
}

// same as KeyReverseLookupAttr where attr == NODES and hint == ""
func (e *EtcdStore) optimizedNodeReverseLookup(key string) (*[]string, error) {
	_, _, value, found, err := e.retrieveFromEtcd(fmt.Sprintf("%s/%s", _roptimize, key), false, false)
	if !found {
		return &[]string{}, errors.New(fmt.Sprintf("Key NOT Found [%s]", key))
	} else if err != nil {
		return &[]string{}, err
	}
	values := strings.Split(value, _sep)
	return &values, nil
}
<|endoftext|>


<commit_before>package hostedtsdb

import (
	"bytes"
	"encoding/gob"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core/ctypes"

	"github.com/raintank/raintank-metric/msg"
	"github.com/raintank/raintank-metric/schema"
)

const (
	name                 = "rt-hostedtsdb"
	version              = 1
	pluginType           = plugin.PublisherPluginType
	maxMetricsPerPayload = 3000
)

var (
	RemoteUrl *url.URL
	Token     string
)

type HostedtsdbPublisher struct {
}

func NewHostedtsdbPublisher() *HostedtsdbPublisher {
	return &HostedtsdbPublisher{}
}

type WriteQueue struct {
	sync.Mutex
	Metrics   []*schema.MetricData
	QueueFull chan struct{}
}

func (q *WriteQueue) Add(metrics []*schema.MetricData) {
	q.Lock()
	q.Metrics = append(q.Metrics, metrics...)
	if len(q.Metrics) > maxMetricsPerPayload {
		q.QueueFull <- struct{}{}
	}
	q.Unlock()
}

func (q *WriteQueue) Flush() {
	q.Lock()
	if len(q.Metrics) == 0 {
		q.Unlock()
		return
	}
	metrics := make([]*schema.MetricData, len(q.Metrics))
	copy(metrics, q.Metrics)
	q.Metrics = q.Metrics[:0]
	q.Unlock()
	// Write the metrics to our HTTP server.
	log.Printf("writing %d metrics to API", len(metrics))
	id := time.Now().UnixNano()
	body, err := msg.CreateMsg(metrics, id, msg.FormatMetricDataArrayMsgp)
	if err != nil {
		log.Printf("Error: unable to convert metrics to MetricDataArrayMsgp. %s", err)
		return
	}
	sent := false
	for !sent {
		if err = PostData(RemoteUrl, Token, body); err != nil {
			log.Printf("Error: %s", err)
			time.Sleep(time.Second)
		} else {
			sent = true
		}
	}
}

func (q *WriteQueue) Run() {
	ticker := time.NewTicker(time.Second)
	for {
		select {
		case <-ticker.C:
			q.Flush()
		case <-q.QueueFull:
			q.Flush()
		}
	}
}

func NewWriteQueue() *WriteQueue {
	return &WriteQueue{
		Metrics:   make([]*schema.MetricData, 0),
		QueueFull: make(chan struct{}),
	}
}

type EventQueue struct {
}

func (e *EventQueue) Add(m *plugin.PluginMetricType) error {
	return nil
}

func NewEventQueue() *EventQueue {
	return &EventQueue{}
}

var writeQueue *WriteQueue

var eventQueue *EventQueue

func init() {
	writeQueue = NewWriteQueue()
	go writeQueue.Run()
	eventQueue = NewEventQueue()
	//go eventQueue.Run()
}

func (f *HostedtsdbPublisher) Publish(contentType string, content []byte, config map[string]ctypes.ConfigValue) error {
	log.Println("Publishing started")
	var metrics []plugin.PluginMetricType

	switch contentType {
	case plugin.SnapGOBContentType:
		dec := gob.NewDecoder(bytes.NewBuffer(content))
		if err := dec.Decode(&metrics); err != nil {
			log.Printf("Error decoding: error=%v content=%v", err, content)
			return err
		}
	default:
		log.Printf("Error unknown content type '%v'", contentType)
		return errors.New(fmt.Sprintf("Unknown content type '%s'", contentType))
	}

	log.Printf("publishing %v metrics to %v", len(metrics), config)

	// set the RemoteURL and Token when the first metrics is recieved.
	var err error
	if RemoteUrl == nil {
		RemoteUrl, err = url.Parse(config["raintank_tsdb_url"].(ctypes.ConfigValueStr).Value)
		if err != nil {
			return err
		}
	}
	if Token == "" {
		Token = config["raintank_api_key"].(ctypes.ConfigValueStr).Value
	}
	//-----------------

	interval := config["interval"].(ctypes.ConfigValueInt).Value
	orgId := config["orgId"].(ctypes.ConfigValueInt).Value

	metricsArray := make([]*schema.MetricData, len(metrics))
	for i, m := range metrics {
		tags := make([]string, 0)
		targetType := "gauge"
		unit := ""
		for k, v := range m.Tags() {
			switch k {
			case "targetType":
				targetType = v
			case "unit":
				unit = v
			default:
				tags = append(tags, fmt.Sprintf("%s:%s", k, v))
			}
		}
		var value float64
		rawData := m.Data()
		switch rawData.(type) {
		case int:
			value = float64(rawData.(int))
		case int8:
			value = float64(rawData.(int8))
		case int16:
			value = float64(rawData.(int16))
		case int32:
			value = float64(rawData.(int32))
		case int64:
			value = float64(rawData.(int64))
		case uint8:
			value = float64(rawData.(uint8))
		case uint16:
			value = float64(rawData.(uint16))
		case uint32:
			value = float64(rawData.(uint32))
		case uint64:
			value = float64(rawData.(uint64))
		case float32:
			value = float64(rawData.(float32))
		case float64:
			value = rawData.(float64)
		case string:
			//payload is an event.
			return eventQueue.Add(&m)
		default:
			return errors.New("unknown data type")
		}

		metricsArray[i] = &schema.MetricData{
			OrgId:      orgId,
			Name:       strings.Join(m.Namespace(), "."),
			Interval:   interval,
			Value:      value,
			Time:       m.Timestamp().Unix(),
			TargetType: targetType,
			Unit:       unit,
			Tags:       tags,
		}
		metricsArray[i].SetId()
	}
	writeQueue.Add(metricsArray)

	return nil
}

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(
		name,
		version,
		pluginType,
		[]string{plugin.SnapGOBContentType},
		[]string{plugin.SnapGOBContentType},
		plugin.Exclusive(true),
	)
}

func (f *HostedtsdbPublisher) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	c := cpolicy.New()
	rule, _ := cpolicy.NewStringRule("raintank_tsdb_url", true)
	rule2, _ := cpolicy.NewStringRule("raintank_api_key", true)
	rule3, _ := cpolicy.NewIntegerRule("interval", true)
	rule4, _ := cpolicy.NewIntegerRule("orgId", false, 0)

	p := cpolicy.NewPolicyNode()
	p.Add(rule)
	p.Add(rule2)
	p.Add(rule3)
	p.Add(rule4)
	c.Add([]string{""}, p)
	return c, nil
}

func handleErr(e error) {
	if e != nil {
		panic(e)
	}
}

func PostData(remoteUrl *url.URL, token string, body []byte) error {
	req, err := http.NewRequest("POST", remoteUrl.String(), bytes.NewBuffer(body))
	req.Header.Set("Content-Type", "rt-metric-binary")
	req.Header.Set("Authorization", "Bearer "+token)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	respBody, _ := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("Posting data failed. %d - %s", resp.StatusCode, string(respBody))
	}
	return nil
}
<commit_msg>add event support to snap-publisher<commit_after>package hostedtsdb

import (
	"bytes"
	"encoding/gob"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core/ctypes"

	"github.com/raintank/raintank-metric/msg"
	"github.com/raintank/raintank-metric/schema"
)

const (
	name                 = "rt-hostedtsdb"
	version              = 1
	pluginType           = plugin.PublisherPluginType
	maxMetricsPerPayload = 3000
)

var (
	RemoteUrl *url.URL
	Token     string
)

type HostedtsdbPublisher struct {
}

func NewHostedtsdbPublisher() *HostedtsdbPublisher {
	return &HostedtsdbPublisher{}
}

type WriteQueue struct {
	sync.Mutex
	Metrics   []*schema.MetricData
	QueueFull chan struct{}
}

func (q *WriteQueue) Add(metrics []*schema.MetricData) {
	q.Lock()
	q.Metrics = append(q.Metrics, metrics...)
	if len(q.Metrics) > maxMetricsPerPayload {
		q.QueueFull <- struct{}{}
	}
	q.Unlock()
}

func (q *WriteQueue) Flush() {
	q.Lock()
	if len(q.Metrics) == 0 {
		q.Unlock()
		return
	}
	metrics := make([]*schema.MetricData, len(q.Metrics))
	copy(metrics, q.Metrics)
	q.Metrics = q.Metrics[:0]
	q.Unlock()
	// Write the metrics to our HTTP server.
	log.Printf("writing %d metrics to API", len(metrics))
	id := time.Now().UnixNano()
	body, err := msg.CreateMsg(metrics, id, msg.FormatMetricDataArrayMsgp)
	if err != nil {
		log.Printf("Error: unable to convert metrics to MetricDataArrayMsgp. %s", err)
		return
	}
	sent := false
	for !sent {
		if err = PostData("metrics", Token, body); err != nil {
			log.Printf("Error: %s", err)
			time.Sleep(time.Second)
		} else {
			sent = true
		}
	}
}

func (q *WriteQueue) Run() {
	ticker := time.NewTicker(time.Second)
	for {
		select {
		case <-ticker.C:
			q.Flush()
		case <-q.QueueFull:
			q.Flush()
		}
	}
}

func NewWriteQueue() *WriteQueue {
	return &WriteQueue{
		Metrics:   make([]*schema.MetricData, 0),
		QueueFull: make(chan struct{}),
	}
}

var writeQueue *WriteQueue

func init() {
	writeQueue = NewWriteQueue()
	go writeQueue.Run()
}

func (f *HostedtsdbPublisher) Publish(contentType string, content []byte, config map[string]ctypes.ConfigValue) error {
	log.Println("Publishing started")
	var metrics []plugin.PluginMetricType

	switch contentType {
	case plugin.SnapGOBContentType:
		dec := gob.NewDecoder(bytes.NewBuffer(content))
		if err := dec.Decode(&metrics); err != nil {
			log.Printf("Error decoding: error=%v content=%v", err, content)
			return err
		}
	default:
		log.Printf("Error unknown content type '%v'", contentType)
		return errors.New(fmt.Sprintf("Unknown content type '%s'", contentType))
	}

	log.Printf("publishing %d metrics to %v", len(metrics), config)

	// set the RemoteURL and Token when the first metrics is recieved.
	var err error
	if RemoteUrl == nil {
		remote := config["raintank_tsdb_url"].(ctypes.ConfigValueStr).Value
		if !strings.HasSuffix(remote, "/") {
			remote += "/"
		}
		RemoteUrl, err = url.Parse(remote)
		if err != nil {
			return err
		}

	}
	if Token == "" {
		Token = config["raintank_api_key"].(ctypes.ConfigValueStr).Value
	}
	//-----------------

	interval := config["interval"].(ctypes.ConfigValueInt).Value
	orgId := config["orgId"].(ctypes.ConfigValueInt).Value

	metricsArray := make([]*schema.MetricData, len(metrics))
	for i, m := range metrics {
		var value float64
		rawData := m.Data()
		switch rawData.(type) {
		case string:
			//payload is an event.
			go sendEvent(int64(orgId), &m)
			continue
		case int:
			value = float64(rawData.(int))
		case int8:
			value = float64(rawData.(int8))
		case int16:
			value = float64(rawData.(int16))
		case int32:
			value = float64(rawData.(int32))
		case int64:
			value = float64(rawData.(int64))
		case uint8:
			value = float64(rawData.(uint8))
		case uint16:
			value = float64(rawData.(uint16))
		case uint32:
			value = float64(rawData.(uint32))
		case uint64:
			value = float64(rawData.(uint64))
		case float32:
			value = float64(rawData.(float32))
		case float64:
			value = rawData.(float64)
		default:
			return errors.New("unknown data type")
		}

		tags := make([]string, 0)
		targetType := "gauge"
		unit := ""
		for k, v := range m.Tags() {
			switch k {
			case "targetType":
				targetType = v
			case "unit":
				unit = v
			default:
				tags = append(tags, fmt.Sprintf("%s:%s", k, v))
			}
		}

		metricsArray[i] = &schema.MetricData{
			OrgId:      orgId,
			Name:       strings.Join(m.Namespace(), "."),
			Interval:   interval,
			Value:      value,
			Time:       m.Timestamp().Unix(),
			TargetType: targetType,
			Unit:       unit,
			Tags:       tags,
		}
		metricsArray[i].SetId()
	}
	writeQueue.Add(metricsArray)

	return nil
}

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(
		name,
		version,
		pluginType,
		[]string{plugin.SnapGOBContentType},
		[]string{plugin.SnapGOBContentType},
		plugin.Exclusive(true),
	)
}

func (f *HostedtsdbPublisher) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	c := cpolicy.New()
	rule, _ := cpolicy.NewStringRule("raintank_tsdb_url", true)
	rule2, _ := cpolicy.NewStringRule("raintank_api_key", true)
	rule3, _ := cpolicy.NewIntegerRule("interval", true)
	rule4, _ := cpolicy.NewIntegerRule("orgId", false, 0)

	p := cpolicy.NewPolicyNode()
	p.Add(rule)
	p.Add(rule2)
	p.Add(rule3)
	p.Add(rule4)
	c.Add([]string{""}, p)
	return c, nil
}

func handleErr(e error) {
	if e != nil {
		panic(e)
	}
}

func PostData(path, token string, body []byte) error {
	u := RemoteUrl.String() + path
	req, err := http.NewRequest("POST", u, bytes.NewBuffer(body))
	req.Header.Set("Content-Type", "rt-metric-binary")
	req.Header.Set("Authorization", "Bearer "+token)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	respBody, _ := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("Posting data failed. %d - %s", resp.StatusCode, string(respBody))
	}
	return nil
}

func sendEvent(orgId int64, m *plugin.PluginMetricType) {
	ns := m.Namespace()
	if len(ns) != 4 {
		log.Printf("Error: invalid event metric. Expected namesapce to be 4 fields.")
		return
	}
	if ns[0] != "worldping" || ns[1] != "event" {
		log.Printf("Error: invalid event metrics.  Metrics hould begin with 'worldping.event'")
		return
	}

	id := time.Now().UnixNano()
	event := &schema.ProbeEvent{
		OrgId:     orgId,
		EventType: ns[2],
		Severity:  ns[3],
		Source:    m.Source(),
		Timestamp: id / int64(time.Millisecond),
		Message:   m.Data().(string),
		Tags:      m.Tags(),
	}

	body, err := msg.CreateProbeEventMsg(event, id, msg.FormatProbeEventMsgp)
	if err != nil {
		log.Printf("Error: unable to convert event to ProbeEventMsgp. %s", err)
		return
	}
	sent := false
	for !sent {
		if err = PostData("events", Token, body); err != nil {
			log.Printf("Error: %s", err)
			time.Sleep(time.Second)
		} else {
			sent = true
		}
	}
}
<|endoftext|>


<commit_before>package a

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"lutils/httplib"
	"lutils/logs"
	"lutils/mahonia"
	"reflect"
	"sort"
	"strings"
	"time"
)

//公共请求参数
type requestParams struct {
	AppId        string `ali:"app_id"`
	Method       string `ali:"method"`
	Format       string `ali:"format"`
	Charset      string `ali:"charset"`
	SignType     string `ali:"sign_type"`
	Sign         string `ali:"sign"`
	TimeStamp    string `ali:"timestamp"`
	Version      string `ali:"version"`
	AppAuthToken string `ali:"app_auth_token"`
	BizContent   string `ali:"biz_content"`
	AuthToken    string `ali:"auth_token"`
	Code         string `ali:"code"`
	GrantType    string `ali:"grant_type"`
}

func (params *requestParams) valid() error {
	if len(params.AppId) == 0 {
		return ErrAppIdNil
	}

	if len(params.Method) == 0 {
		return errors.New("method 不能为空")
	}

	if len(params.Format) == 0 {
		params.Format = "JSON"
	}

	if len(params.Format) > 0 && params.Format != "JSON" {
		return errors.New("format 仅支持JSON")
	}

	if len(params.Charset) == 0 {
		params.Charset = "GBK"
	}

	if len(params.Charset) > 0 && strings.ToUpper(params.Charset) != "GBK" {
		return errors.New("charset 目前仅支持GBK")
	}

	if len(params.SignType) == 0 {
		params.SignType = "RSA"
	}

	if len(params.SignType) > 0 && params.SignType != "RSA" {
		return errors.New("sign_type 仅支持RSA")
	}

	if len(params.TimeStamp) == 0 {
		params.TimeStamp = time.Now().Format("2006-01-02 15:04:05")
	}

	if len(params.TimeStamp) > 0 {
		if _, err := time.Parse("2006-01-02 15:04:05", params.TimeStamp); err != nil {
			return errors.New("timestamp 格式\"yyyy-MM-dd HH:mm:ss\"")
		}
	}

	if len(params.Version) == 0 {
		params.Version = "1.0"
	} else if params.Version != "1.0" {
		return errors.New("version 固定为：1.0")
	}

	return nil
}

type bizInterface interface {
	valid() error
}

type responseInterface interface{}

type Response struct {
	Code    string `json:"code,omitempty"`
	Msg     string `json:"msg,omitempty"`
	SubCode string `json:"sub_code,omitempty"`
	SubMsg  string `json:"sub_msg,omitempty"`
}

type ApiHander interface {
	apiMethod() string
	apiName() string
}

var (
	apiRegistry map[string]AlipayApi
)

func init() {
	apiRegistry = map[string]AlipayApi{}
}

func registerApi(handler ApiHander) {
	apiRegistry[handler.apiMethod()] = AlipayApi{
		apiname:   handler.apiName,
		apimethod: handler.apiMethod,
	}
}

func GetApi(method string) AlipayApi {
	return apiRegistry[method]
}

func GetSupportApis() string {
	lst := "\n=====================SUPPORTED ALIPAY API LIST=====================\n"
	for _, v := range apiRegistry {
		lst += "====[" + v.apiname() + ":" + v.apimethod() + "]====" + "\n"
	}
	return lst
}

type AlipayApi struct {
	params    requestParams
	apiname   func() string
	apimethod func() string
}

func (a *AlipayApi) SetAppId(app_id string) error {
	a.params.AppId = app_id
	if len(a.params.AppId) == 0 {
		return ErrAppIdNil
	}

	if _, ok := secretLst[a.params.AppId]; !ok {
		return ErrSecretNil
	}
	return nil
}

func (a *AlipayApi) SetBizContent(biz bizInterface) error {
	//有几个接口要独立处理
	if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_system_oauth_token{}).Name() {
		b := biz.(Biz_alipay_system_oauth_token)
		a.params.Code = b.Code
		a.params.GrantType = b.GrantType
		a.params.BizContent = ""
		return nil
	} else if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_user_info_share{}).Name() {
		b := biz.(Biz_alipay_user_info_share)
		a.params.AuthToken = b.AuthToken
		a.params.BizContent = ""
		return nil
	}

	if v, err := a.biz_to_string(biz); err != nil {
		return err
	} else {
		a.params.BizContent = v
		return nil
	}
}

func (a *AlipayApi) struct_to_map() map[string]interface{} {
	t := reflect.TypeOf(a.params)
	v := reflect.ValueOf(a.params)

	var data = make(map[string]interface{})
	for i := 0; i < t.NumField(); i++ {
		key := t.Field(i).Name
		value := v.Field(i).Interface()
		tag := t.Field(i).Tag.Get("ali")
		if tag != "" {
			if strings.Contains(tag, ",") {
				ps := strings.Split(tag, ",")
				key = ps[0]
			} else {
				key = tag
			}
		}
		data[key] = value
	}
	return data
}

func (a *AlipayApi) utf_to_gbk(utf string) (gbk string) {
	gbk_enc := mahonia.NewEncoder("GBK")
	return gbk_enc.ConvertString(utf)
}

func (a *AlipayApi) gbk_to_utf(gbk string) (utf string) {
	enc := mahonia.NewDecoder("GBK")
	return enc.ConvertString(gbk)
}

func (a *AlipayApi) map_to_string(m map[string]interface{}) string {
	//对key进行升序排序.
	sorted_keys := make([]string, 0)
	for k, _ := range m {
		sorted_keys = append(sorted_keys, k)
	}
	sort.Strings(sorted_keys)

	//对key=value的键值对用&连接起来，略过空值
	var signStrings string
	for _, k := range sorted_keys {
		value := fmt.Sprintf("%v", m[k])
		if value != "" {
			signStrings = signStrings + k + "=" + value + "&"
		}
	}
	if len(signStrings) == 0 {
		return ""
	} else {
		signStrings = signStrings[:len(signStrings)-1]
	}
	return signStrings
}

func (a *AlipayApi) biz_to_string(b bizInterface) (string, error) {
	if err := b.valid(); err != nil {
		return "", err
	}
	content := ""
	if v, err := json.Marshal(&b); err != nil {
		return "", err
	} else {
		content = string(v)
	}

	temp_map := map[string]interface{}{
		"biz": content,
	}

	if v, err := json.Marshal(&temp_map); err != nil {
		return "", err
	} else {
		content = string(v)
	}
	return content[8 : len(content)-2], nil
}

func (a *AlipayApi) sign(c string) (sign string, err error) {
	//签名
	s := getSecret(a.params.AppId)
	block, _ := pem.Decode(s.PrivRSA)
	if block == nil {
		return "", errors.New("privateKey error!")
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}
	_t := crypto.SHA1.New()
	_t.Write([]byte(c))
	digest := _t.Sum(nil)
	rsa_data, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA1, digest)
	if err != nil {
		return "", err
	}
	result := string(rsa_data)
	result = base64.StdEncoding.EncodeToString(rsa_data)
	return result, nil
}

func (a *AlipayApi) verifySign(s, origin_sign, method_key string) bool {
	sign_start_index := strings.Index(s, ",\"sign\"")
	if sign_start_index == -1 {
		return false
	}
	tobe_signed := s[4+len(method_key) : sign_start_index]
	block, _ := pem.Decode(getSecret(a.params.AppId).AliPubRSA)
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		fmt.Printf("Failed to parse RSA public key: %s\n", err)
		return false
	}
	rsaPub, _ := pub.(*rsa.PublicKey)
	t := sha1.New()
	io.WriteString(t, tobe_signed)
	digest := t.Sum(nil)
	data, err := base64.StdEncoding.DecodeString(origin_sign)
	if err != nil {
		fmt.Println("DecodeString sig error, reason: ", err)
		return false
	}
	err = rsa.VerifyPKCS1v15(rsaPub, crypto.SHA1, digest, data)
	if err != nil {
		fmt.Println("Verify sig error, reason: ", err)
		return false
	}

	return true
}

func (a *AlipayApi) request(m map[string]interface{}) (string, error) {
	url_link := "https://openapi.alipay.com/gateway.do"
	if conf.SandBoxEnable {
		url_link = "https://openapi.alipaydev.com/gateway.do"
	}
	http_request := httplib.Post(url_link)
	tmp_string := ""
	for k, _ := range m {
		value := fmt.Sprintf("%v", m[k])
		if value != "" {
			http_request.Param(k, value)
			tmp_string = tmp_string + k + "=" + value + "\t"
		}
	}

	logs.DEBUG(fmt.Sprintf("==[请求参数]==[%s]", tmp_string))
	var string_result string
	if v, err := http_request.String(); err != nil {
		return "", err
	} else {
		string_result = v

	}
	return string_result, nil
}

func (a *AlipayApi) Run(resp responseInterface) error {
	defer logs.DEBUG("=====================ALIPAY REQUEST END=====================")
	logs.DEBUG("=====================ALIPAY REQUEST START=====================")
	logs.DEBUG(fmt.Sprintf("==[沙盒模式]==[%v]", conf.SandBoxEnable))
	logs.DEBUG(fmt.Sprintf("==[调用方法]==[%s]:[%s]", a.apiname(), a.apimethod()))
	a.params.Method = a.apimethod()
	if err := a.params.valid(); err != nil {
		fmt.Println(err.Error())
		return err
	}

	rv := reflect.ValueOf(resp)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		if rv.Type() == nil {
			return errors.New("responseInterface is nil")
		}

		if rv.Type().Kind() != reflect.Ptr {
			return errors.New("responseInterface is non-pointer" + rv.Type().String())
		}
		return errors.New("responseInterface is nil" + rv.Type().String())
	}

	m := a.struct_to_map()
	//做请求参数的签名
	__sign := ""
	tobe_sign := a.map_to_string(m)
	logs.DEBUG(fmt.Sprintf("==[准备签名]==[%s]", tobe_sign))
	if v, err := a.sign(tobe_sign); err != nil {
		return err
	} else if len(v) == 0 {
		return ErrSign
	} else {
		__sign = v
	}
	logs.DEBUG(fmt.Sprintf("==[签名结果]==[%s]", __sign))
	m["sign"] = __sign
	//准备请求
	result_string := ""
	if v, err := a.request(m); err != nil {
		return err
	} else {
		result_string = v
		logs.DEBUG(fmt.Sprintf("==[响应结果]==[GBK 编码]:[%s]", result_string))
	}

	//把 method 转换为 key
	method_key := strings.Replace(a.params.Method, ".", "_", -1)
	method_key += "_response"

	result_string = strings.Replace(result_string, "error_response", method_key, -1)
	//解析结果
	resp_map := map[string]interface{}{}
	if err := json.Unmarshal([]byte(result_string), &resp_map); err != nil {
		return err
	}
	//原始签名
	if v, ok := resp_map["sign"].(string); ok && len(v) > 0 {
		if pass := a.verifySign(result_string, v, method_key); !pass {
			return ErrVerifySign
		}
	}

	//转码
	result_string = a.gbk_to_utf(result_string)
	logs.DEBUG(fmt.Sprintf("==[响应结果]==[UTF 编码]:[%s]", result_string))

	if err := json.Unmarshal([]byte(result_string), &resp_map); err != nil {
		return err
	}
	//把需要的内容再次换成string
	if v, err := json.Marshal(resp_map[method_key]); err != nil {
		return err
	} else {
		result_string = string(v)
	}

	if err := json.Unmarshal([]byte(result_string), resp); err != nil {
		return err
	}
	return nil
}
<commit_msg>修改 userinfo 接口的 bug<commit_after>package a

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"lutils/httplib"
	"lutils/logs"
	"lutils/mahonia"
	"reflect"
	"sort"
	"strings"
	"time"
)

//公共请求参数
type requestParams struct {
	AppId        string `ali:"app_id"`
	Method       string `ali:"method"`
	Format       string `ali:"format"`
	Charset      string `ali:"charset"`
	SignType     string `ali:"sign_type"`
	Sign         string `ali:"sign"`
	TimeStamp    string `ali:"timestamp"`
	Version      string `ali:"version"`
	AppAuthToken string `ali:"app_auth_token"`
	BizContent   string `ali:"biz_content"`
	AuthToken    string `ali:"auth_token"`
	Code         string `ali:"code"`
	GrantType    string `ali:"grant_type"`
}

func (params *requestParams) valid() error {
	if len(params.AppId) == 0 {
		return ErrAppIdNil
	}

	if len(params.Method) == 0 {
		return errors.New("method 不能为空")
	}

	if len(params.Format) == 0 {
		params.Format = "JSON"
	}

	if len(params.Format) > 0 && params.Format != "JSON" {
		return errors.New("format 仅支持JSON")
	}

	if len(params.Charset) == 0 {
		params.Charset = "GBK"
	}

	if len(params.Charset) > 0 && strings.ToUpper(params.Charset) != "GBK" {
		return errors.New("charset 目前仅支持GBK")
	}

	if len(params.SignType) == 0 {
		params.SignType = "RSA"
	}

	if len(params.SignType) > 0 && params.SignType != "RSA" {
		return errors.New("sign_type 仅支持RSA")
	}

	if len(params.TimeStamp) == 0 {
		params.TimeStamp = time.Now().Format("2006-01-02 15:04:05")
	}

	if len(params.TimeStamp) > 0 {
		if _, err := time.Parse("2006-01-02 15:04:05", params.TimeStamp); err != nil {
			return errors.New("timestamp 格式\"yyyy-MM-dd HH:mm:ss\"")
		}
	}

	if len(params.Version) == 0 {
		params.Version = "1.0"
	} else if params.Version != "1.0" {
		return errors.New("version 固定为：1.0")
	}

	return nil
}

type bizInterface interface {
	valid() error
}

type responseInterface interface{}

type Response struct {
	Code    string `json:"code,omitempty"`
	Msg     string `json:"msg,omitempty"`
	SubCode string `json:"sub_code,omitempty"`
	SubMsg  string `json:"sub_msg,omitempty"`
}

type ApiHander interface {
	apiMethod() string
	apiName() string
}

var (
	apiRegistry map[string]AlipayApi
)

func init() {
	apiRegistry = map[string]AlipayApi{}
}

func registerApi(handler ApiHander) {
	apiRegistry[handler.apiMethod()] = AlipayApi{
		apiname:   handler.apiName,
		apimethod: handler.apiMethod,
	}
}

func GetApi(method string) AlipayApi {
	return apiRegistry[method]
}

func GetSupportApis() string {
	lst := "\n=====================SUPPORTED ALIPAY API LIST=====================\n"
	for _, v := range apiRegistry {
		lst += "====[" + v.apiname() + ":" + v.apimethod() + "]====" + "\n"
	}
	return lst
}

type AlipayApi struct {
	params    requestParams
	apiname   func() string
	apimethod func() string
}

func (a *AlipayApi) SetAppId(app_id string) error {
	a.params.AppId = app_id
	if len(a.params.AppId) == 0 {
		return ErrAppIdNil
	}

	if _, ok := secretLst[a.params.AppId]; !ok {
		return ErrSecretNil
	}
	return nil
}

func (a *AlipayApi) SetBizContent(biz bizInterface) error {
	//有几个接口要独立处理
	if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_system_oauth_token{}).Name() {
		b := biz.(Biz_alipay_system_oauth_token)
		a.params.Code = b.Code
		a.params.GrantType = b.GrantType
		a.params.BizContent = ""
		return nil
	} else if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_user_info_share{}).Name() {
		b := biz.(Biz_alipay_user_info_share)
		a.params.AuthToken = b.AuthToken
		a.params.BizContent = ""
		return nil
	} else if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_user_userinfo_share{}).Name() {
		b := biz.(Biz_alipay_user_userinfo_share)
		a.params.AuthToken = b.AuthToken
		a.params.BizContent = ""
		return nil
	}

	if v, err := a.biz_to_string(biz); err != nil {
		return err
	} else {
		a.params.BizContent = v
		return nil
	}
}

func (a *AlipayApi) struct_to_map() map[string]interface{} {
	t := reflect.TypeOf(a.params)
	v := reflect.ValueOf(a.params)

	var data = make(map[string]interface{})
	for i := 0; i < t.NumField(); i++ {
		key := t.Field(i).Name
		value := v.Field(i).Interface()
		tag := t.Field(i).Tag.Get("ali")
		if tag != "" {
			if strings.Contains(tag, ",") {
				ps := strings.Split(tag, ",")
				key = ps[0]
			} else {
				key = tag
			}
		}
		data[key] = value
	}
	return data
}

func (a *AlipayApi) utf_to_gbk(utf string) (gbk string) {
	gbk_enc := mahonia.NewEncoder("GBK")
	return gbk_enc.ConvertString(utf)
}

func (a *AlipayApi) gbk_to_utf(gbk string) (utf string) {
	enc := mahonia.NewDecoder("GBK")
	return enc.ConvertString(gbk)
}

func (a *AlipayApi) map_to_string(m map[string]interface{}) string {
	//对key进行升序排序.
	sorted_keys := make([]string, 0)
	for k, _ := range m {
		sorted_keys = append(sorted_keys, k)
	}
	sort.Strings(sorted_keys)

	//对key=value的键值对用&连接起来，略过空值
	var signStrings string
	for _, k := range sorted_keys {
		value := fmt.Sprintf("%v", m[k])
		if value != "" {
			signStrings = signStrings + k + "=" + value + "&"
		}
	}
	if len(signStrings) == 0 {
		return ""
	} else {
		signStrings = signStrings[:len(signStrings)-1]
	}
	return signStrings
}

func (a *AlipayApi) biz_to_string(b bizInterface) (string, error) {
	if err := b.valid(); err != nil {
		return "", err
	}
	content := ""
	if v, err := json.Marshal(&b); err != nil {
		return "", err
	} else {
		content = string(v)
	}

	temp_map := map[string]interface{}{
		"biz": content,
	}

	if v, err := json.Marshal(&temp_map); err != nil {
		return "", err
	} else {
		content = string(v)
	}
	return content[8 : len(content)-2], nil
}

func (a *AlipayApi) sign(c string) (sign string, err error) {
	//签名
	s := getSecret(a.params.AppId)
	block, _ := pem.Decode(s.PrivRSA)
	if block == nil {
		return "", errors.New("privateKey error!")
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}
	_t := crypto.SHA1.New()
	_t.Write([]byte(c))
	digest := _t.Sum(nil)
	rsa_data, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA1, digest)
	if err != nil {
		return "", err
	}
	result := string(rsa_data)
	result = base64.StdEncoding.EncodeToString(rsa_data)
	return result, nil
}

func (a *AlipayApi) verifySign(s, origin_sign, method_key string) bool {
	sign_start_index := strings.Index(s, ",\"sign\"")
	if sign_start_index == -1 {
		return false
	}
	tobe_signed := s[4+len(method_key) : sign_start_index]
	block, _ := pem.Decode(getSecret(a.params.AppId).AliPubRSA)
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		fmt.Printf("Failed to parse RSA public key: %s\n", err)
		return false
	}
	rsaPub, _ := pub.(*rsa.PublicKey)
	t := sha1.New()
	io.WriteString(t, tobe_signed)
	digest := t.Sum(nil)
	data, err := base64.StdEncoding.DecodeString(origin_sign)
	if err != nil {
		fmt.Println("DecodeString sig error, reason: ", err)
		return false
	}
	err = rsa.VerifyPKCS1v15(rsaPub, crypto.SHA1, digest, data)
	if err != nil {
		fmt.Println("Verify sig error, reason: ", err)
		return false
	}

	return true
}

func (a *AlipayApi) request(m map[string]interface{}) (string, error) {
	url_link := "https://openapi.alipay.com/gateway.do"
	if conf.SandBoxEnable {
		url_link = "https://openapi.alipaydev.com/gateway.do"
	}
	http_request := httplib.Post(url_link)
	tmp_string := ""
	for k, _ := range m {
		value := fmt.Sprintf("%v", m[k])
		if value != "" {
			http_request.Param(k, value)
			tmp_string = tmp_string + k + "=" + value + "\t"
		}
	}

	logs.DEBUG(fmt.Sprintf("==[请求参数]==[%s]", tmp_string))
	var string_result string
	if v, err := http_request.String(); err != nil {
		return "", err
	} else {
		string_result = v

	}
	return string_result, nil
}

func (a *AlipayApi) Run(resp responseInterface) error {
	defer logs.DEBUG("=====================ALIPAY REQUEST END=====================")
	logs.DEBUG("=====================ALIPAY REQUEST START=====================")
	logs.DEBUG(fmt.Sprintf("==[沙盒模式]==[%v]", conf.SandBoxEnable))
	logs.DEBUG(fmt.Sprintf("==[调用方法]==[%s]:[%s]", a.apiname(), a.apimethod()))
	a.params.Method = a.apimethod()
	if err := a.params.valid(); err != nil {
		fmt.Println(err.Error())
		return err
	}

	rv := reflect.ValueOf(resp)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		if rv.Type() == nil {
			return errors.New("responseInterface is nil")
		}

		if rv.Type().Kind() != reflect.Ptr {
			return errors.New("responseInterface is non-pointer" + rv.Type().String())
		}
		return errors.New("responseInterface is nil" + rv.Type().String())
	}

	m := a.struct_to_map()
	//做请求参数的签名
	__sign := ""
	tobe_sign := a.map_to_string(m)
	logs.DEBUG(fmt.Sprintf("==[准备签名]==[%s]", tobe_sign))
	if v, err := a.sign(tobe_sign); err != nil {
		return err
	} else if len(v) == 0 {
		return ErrSign
	} else {
		__sign = v
	}
	logs.DEBUG(fmt.Sprintf("==[签名结果]==[%s]", __sign))
	m["sign"] = __sign
	//准备请求
	result_string := ""
	if v, err := a.request(m); err != nil {
		return err
	} else {
		result_string = v
		logs.DEBUG(fmt.Sprintf("==[响应结果]==[GBK 编码]:[%s]", result_string))
	}

	//把 method 转换为 key
	method_key := strings.Replace(a.params.Method, ".", "_", -1)
	method_key += "_response"

	result_string = strings.Replace(result_string, "error_response", method_key, -1)
	//解析结果
	resp_map := map[string]interface{}{}
	if err := json.Unmarshal([]byte(result_string), &resp_map); err != nil {
		return err
	}
	//原始签名
	if v, ok := resp_map["sign"].(string); ok && len(v) > 0 {
		if pass := a.verifySign(result_string, v, method_key); !pass {
			return ErrVerifySign
		}
	}

	//转码
	result_string = a.gbk_to_utf(result_string)
	logs.DEBUG(fmt.Sprintf("==[响应结果]==[UTF 编码]:[%s]", result_string))

	if err := json.Unmarshal([]byte(result_string), &resp_map); err != nil {
		return err
	}
	//把需要的内容再次换成string
	if v, err := json.Marshal(resp_map[method_key]); err != nil {
		return err
	} else {
		result_string = string(v)
	}

	if err := json.Unmarshal([]byte(result_string), resp); err != nil {
		return err
	}
	return nil
}
<|endoftext|>