<commit_before>package task

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"neon/build"
	"os"
	"path/filepath"
	"reflect"
	"strings"
)

func init() {
	build.AddTask(build.TaskDesc{
		Name: "untar",
		Func: Untar,
		Args: reflect.TypeOf(UntarArgs{}),
		Help: `Expand a tar file in a directory.

Arguments:

- untar: the tar file to expand.
- todir: the destination directory.

Examples:

    # untar foo.tar to build directory
    - untar: "foo.tar"
      todir: "build"

Notes:

- If archive filename ends with gz (with a name such as foo.tar.gz or foo.tgz)
  the tar archive is uncompressed with gzip.`,
	})
}

type UntarArgs struct {
	Untar string `file`
	Todir string `file`
}

func Untar(context *build.Context, args interface{}) error {
	params := args.(UntarArgs)
	context.Message("Untarring archive '%s' to directory '%s'...", params.Untar, params.Todir)
	err := UntarFile(params.Untar, params.Todir)
	if err != nil {
		return fmt.Errorf("expanding archive: %v", err)
	}
	return nil
}

// Untar given file to a directory
func UntarFile(file, dir string) error {
	reader, err := os.Open(file)
	if err != nil {
		return fmt.Errorf("opening source tar file %s: %v", file, err)
	}
	defer reader.Close()
	var tarReader *tar.Reader
	if strings.HasSuffix(file, ".gz") || strings.HasSuffix(file, ".tgz") {
		gzipReader, err := gzip.NewReader(reader)
		defer gzipReader.Close()
		if err != nil {
			return fmt.Errorf("unzipping tar file: %v", err)
		}
		tarReader = tar.NewReader(gzipReader)
	} else {
		tarReader = tar.NewReader(reader)
	}
	for {
		header, err := tarReader.Next()
		switch {
		case err == io.EOF:
			return nil
		case err != nil:
			return err
		case header == nil:
			continue
		}
		target := filepath.Join(dir, header.Name)
		if header.Typeflag == tar.TypeReg {
			destination := filepath.Dir(target)
			if _, err := os.Stat(destination); err != nil {
				if err := os.MkdirAll(destination, 0755); err != nil {
					return fmt.Errorf("creating destination director %s: %v", target, err)
				}
			}
			dest, err := os.OpenFile(target, os.O_CREATE|os.O_RDWR, os.FileMode(header.Mode))
			if err != nil {
				return fmt.Errorf("creating destination file %s: %v", target, err)
			}
			if _, err := io.Copy(dest, tarReader); err != nil {
				return fmt.Errorf("copying to destination file %s: %v", target, err)
			}
			dest.Close()
		}
	}
}
<commit_msg>Rewriting untar task and test<commit_after>package task

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"neon/build"
	"os"
	"path/filepath"
	"reflect"
	"strings"
)

func init() {
	build.AddTask(build.TaskDesc{
		Name: "untar",
		Func: Untar,
		Args: reflect.TypeOf(UntarArgs{}),
		Help: `Expand a tar file in a directory.

Arguments:

- untar: the tar file to expand (string, file).
- todir: the destination directory (string, file).

Examples:

    # untar foo.tar to build directory
    - untar: 'foo.tar'
      todir: 'build'

Notes:

- If archive filename ends with .gz (with a name such as foo.tar.gz or foo.tgz)
  the tar archive is uncompressed with gzip.`,
	})
}

type UntarArgs struct {
	Untar string `file`
	Todir string `file`
}

func Untar(context *build.Context, args interface{}) error {
	params := args.(UntarArgs)
	context.Message("Untarring archive '%s' to directory '%s'...", params.Untar, params.Todir)
	err := UntarFile(params.Untar, params.Todir)
	if err != nil {
		return fmt.Errorf("expanding archive: %v", err)
	}
	return nil
}

func UntarFile(file, dir string) error {
	reader, err := os.Open(file)
	if err != nil {
		return fmt.Errorf("opening source tar file %s: %v", file, err)
	}
	defer reader.Close()
	var tarReader *tar.Reader
	if strings.HasSuffix(file, ".gz") || strings.HasSuffix(file, ".tgz") {
		gzipReader, err := gzip.NewReader(reader)
		defer gzipReader.Close()
		if err != nil {
			return fmt.Errorf("unzipping tar file: %v", err)
		}
		tarReader = tar.NewReader(gzipReader)
	} else {
		tarReader = tar.NewReader(reader)
	}
	for {
		header, err := tarReader.Next()
		switch {
		case err == io.EOF:
			return nil
		case err != nil:
			return err
		case header == nil:
			continue
		}
		target := filepath.Join(dir, header.Name)
		if header.Typeflag == tar.TypeReg {
			destination := filepath.Dir(target)
			if _, err := os.Stat(destination); err != nil {
				if err := os.MkdirAll(destination, 0755); err != nil {
					return fmt.Errorf("creating destination director %s: %v", target, err)
				}
			}
			dest, err := os.OpenFile(target, os.O_CREATE|os.O_RDWR, os.FileMode(header.Mode))
			if err != nil {
				return fmt.Errorf("creating destination file %s: %v", target, err)
			}
			if _, err := io.Copy(dest, tarReader); err != nil {
				return fmt.Errorf("copying to destination file %s: %v", target, err)
			}
			dest.Close()
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"sync"
)

func cacl(i int, wg *sync.WaitGroup) {
	t := 0
	for i := 1; i < 1000000; i++ {
		t++
	}

	fmt.Println(i, t)
	wg.Done()
}

func main() {
	wg := sync.WaitGroup{}
	wg.Add(10)
	for i := 0; i < 10; i++ {
		go cacl(i, &wg)
	}

	wg.Wait()
	fmt.Println("End!!!!")
}
<commit_msg>using defer<commit_after>package main

import (
	"fmt"
	"sync"
)

func cacl(i int, wg *sync.WaitGroup) {
	defer wg.Done()
	t := 0
	for i := 1; i < 1000000; i++ {
		t++
	}

	fmt.Println(i, t)
}

func main() {
	wg := sync.WaitGroup{}
	wg.Add(10)
	for i := 0; i < 10; i++ {
		go cacl(i, &wg)
	}

	wg.Wait()
	fmt.Println("End!!!!")
}
<|endoftext|>


<commit_before>package profile

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/persist"
)

// There's a global lock on cpu and memory profiling, because I'm not sure what
// happens if multiple threads call each at the same time. This lock might be
// unnecessary.
var (
	cpuActive bool
	cpuLock   sync.Mutex
	memActive bool
	memLock   sync.Mutex
)

// startCPUProfile starts cpu profiling. An error will be returned if a cpu
// profiler is already running.
func StartCPUProfile(profileDir, identifier string) error {
	// Lock the cpu profile lock so that only one profiler is running at a
	// time.
	cpuLock.Lock()
	if cpuActive {
		cpuLock.Unlock()
		return errors.New("cannot start cpu profilier, a profiler is already running")
	}
	cpuActive = true
	cpuLock.Unlock()

	// Start profiling into the profile dir, using the identifer. The timestamp
	// of the start time of the profiling will be included in the filename.
	cpuProfileFile, err := os.Create(filepath.Join(profileDir, "cpu-profile-"+identifier+"-"+time.Now().Format(time.RFC3339Nano)+".prof"))
	if err != nil {
		return err
	}
	pprof.StartCPUProfile(cpuProfileFile)
	return nil
}

// stopCPUProfile stops cpu profiling.
func StopCPUProfile() {
	cpuLock.Lock()
	if cpuActive {
		pprof.StopCPUProfile()
		cpuActive = false
	}
	cpuLock.Unlock()
}

// saveMemProfile saves the current memory structure of the program. An error
// will be returned if memory profiling is already in progress. Unlike for cpu
// profiling, there is no 'stopMemProfile' call - everything happens at once.
func SaveMemProfile(profileDir, identifier string) error {
	memLock.Lock()
	if memActive {
		memLock.Unlock()
		return errors.New("cannot start memory profiler, a memory profiler is already running")
	}
	memActive = true
	memLock.Unlock()

	// Save the memory profile.
	memFile, err := os.Create(filepath.Join(profileDir, "mem-profile-"+identifier+"-"+time.Now().Format(time.RFC3339Nano)+".prof"))
	if err != nil {
		return err
	}
	pprof.WriteHeapProfile(memFile)

	memLock.Lock()
	memActive = false
	memLock.Unlock()
	return nil
}

// StartContinuousProfiling will continuously print statistics about the cpu
// usage, memory usage, and runtime stats of the program.
func StartContinuousProfile(profileDir string) {
	// Create the folder for all of the profiling results.
	err := os.MkdirAll(profileDir, 0700)
	if err != nil {
		fmt.Println(err)
		return
	}

	// Continuously log statistics about the running Sia application.
	go func() {
		// Create the logger.
		log, err := persist.NewFileLogger(filepath.Join(profileDir, "continuousProfiling.log"))
		if err != nil {
			fmt.Println("Profile logging failed:", err)
			return
		}

		// Collect statistics in an infinite loop.
		sleepTime := time.Second * 20
		for {
			// Sleep for an exponential amount of time each iteration, this
			// keeps the size of the log small while still providing lots of
			// information.
			time.Sleep(sleepTime)
			sleepTime = time.Duration(1.5 * float64(sleepTime))

			var m runtime.MemStats
			runtime.ReadMemStats(&m)
			log.Printf("\n\tGoroutines: %v\n\tAlloc: %v\n\tTotalAlloc: %v\n\tHeapAlloc: %v\n\tHeapSys: %v\n", runtime.NumGoroutine(), m.Alloc, m.TotalAlloc, m.HeapAlloc, m.HeapSys)
		}
	}()
}
<commit_msg>profile flag will grap CPU and Mem profiles<commit_after>package profile

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/persist"
)

// There's a global lock on cpu and memory profiling, because I'm not sure what
// happens if multiple threads call each at the same time. This lock might be
// unnecessary.
var (
	cpuActive bool
	cpuLock   sync.Mutex
	memActive bool
	memLock   sync.Mutex
)

// startCPUProfile starts cpu profiling. An error will be returned if a cpu
// profiler is already running.
func StartCPUProfile(profileDir, identifier string) error {
	// Lock the cpu profile lock so that only one profiler is running at a
	// time.
	cpuLock.Lock()
	if cpuActive {
		cpuLock.Unlock()
		return errors.New("cannot start cpu profilier, a profiler is already running")
	}
	cpuActive = true
	cpuLock.Unlock()

	// Start profiling into the profile dir, using the identifer. The timestamp
	// of the start time of the profiling will be included in the filename.
	cpuProfileFile, err := os.Create(filepath.Join(profileDir, "cpu-profile-"+identifier+"-"+time.Now().Format(time.RFC3339Nano)+".prof"))
	if err != nil {
		return err
	}
	pprof.StartCPUProfile(cpuProfileFile)
	return nil
}

// stopCPUProfile stops cpu profiling.
func StopCPUProfile() {
	cpuLock.Lock()
	if cpuActive {
		pprof.StopCPUProfile()
		cpuActive = false
	}
	cpuLock.Unlock()
}

// saveMemProfile saves the current memory structure of the program. An error
// will be returned if memory profiling is already in progress. Unlike for cpu
// profiling, there is no 'stopMemProfile' call - everything happens at once.
func SaveMemProfile(profileDir, identifier string) error {
	memLock.Lock()
	if memActive {
		memLock.Unlock()
		return errors.New("cannot start memory profiler, a memory profiler is already running")
	}
	memActive = true
	memLock.Unlock()

	// Save the memory profile.
	memFile, err := os.Create(filepath.Join(profileDir, "mem-profile-"+identifier+"-"+time.Now().Format(time.RFC3339Nano)+".prof"))
	if err != nil {
		return err
	}
	pprof.WriteHeapProfile(memFile)

	memLock.Lock()
	memActive = false
	memLock.Unlock()
	return nil
}

// StartContinuousProfiling will continuously print statistics about the cpu
// usage, memory usage, and runtime stats of the program.
func StartContinuousProfile(profileDir string) {
	// Create the folder for all of the profiling results.
	err := os.MkdirAll(profileDir, 0700)
	if err != nil {
		fmt.Println(err)
		return
	}

	// Continuously log statistics about the running Sia application.
	go func() {
		// Create the logger.
		log, err := persist.NewFileLogger(filepath.Join(profileDir, "continuousProfiling.log"))
		if err != nil {
			fmt.Println("Profile logging failed:", err)
			return
		}

		// Collect statistics in an infinite loop.
		sleepTime := time.Second * 20
		for {
			// Sleep for an exponential amount of time each iteration, this
			// keeps the size of the log small while still providing lots of
			// information.
			StopCPUProfile()
			SaveMemProfile(profileDir, "continuousProfilingMem")
			StartCPUProfile(profileDir, "continuousProfilingCPU")
			time.Sleep(sleepTime)
			sleepTime = time.Duration(1.5 * float64(sleepTime))

			var m runtime.MemStats
			runtime.ReadMemStats(&m)
			log.Printf("\n\tGoroutines: %v\n\tAlloc: %v\n\tTotalAlloc: %v\n\tHeapAlloc: %v\n\tHeapSys: %v\n", runtime.NumGoroutine(), m.Alloc, m.TotalAlloc, m.HeapAlloc, m.HeapSys)
		}
	}()
}
<|endoftext|>


<commit_before>package program

import (
	"fmt"
)

// Versioning: [major].[minor].[revision]

const VERSION = "0.1"

type VersionNo struct {
	Major    int
	Minor    int
	Revision int
}

func ParseVersionNumber(verstr string) VersionNo {
	vn := VersionNo{}
	fmt.Sscanf(verstr, "%d.%d.%d", &vn.Major, &vn.Minor, &vn.Revision)
	if vn.Major < 0 {
		vn.Major = 0
	}
	if vn.Minor < 0 {
		vn.Minor = 0
	}
	if vn.Revision < 0 {
		vn.Revision = 0
	}
	return vn
}

func (vn VersionNo) Compare(anotherVersion VersionNo) int {
	if vn.Major < anotherVersion.Major {
		return -1
	}
	if vn.Major > anotherVersion.Major {
		return 1
	}
	if vn.Minor < anotherVersion.Minor {
		return -1
	}
	if vn.Minor > anotherVersion.Minor {
		return 1
	}
	if vn.Revision < anotherVersion.Revision {
		return -1
	}
	if vn.Revision > anotherVersion.Revision {
		return 1
	}
	return 0
}

func (vn VersionNo) String() string {
	if vn.Revision == 0 {
		return fmt.Sprintf("%d.%d", vn.Major, vn.Minor)
	}
	return fmt.Sprintf("%d.%d.%d", vn.Major, vn.Minor, vn.Revision)
}
<commit_msg>Remove unused const.<commit_after>package program

import (
	"fmt"
)

// Versioning: [major].[minor].[revision]

type VersionNo struct {
	Major    int
	Minor    int
	Revision int
}

func ParseVersionNumber(verstr string) VersionNo {
	vn := VersionNo{}
	fmt.Sscanf(verstr, "%d.%d.%d", &vn.Major, &vn.Minor, &vn.Revision)
	if vn.Major < 0 {
		vn.Major = 0
	}
	if vn.Minor < 0 {
		vn.Minor = 0
	}
	if vn.Revision < 0 {
		vn.Revision = 0
	}
	return vn
}

func (vn VersionNo) Compare(anotherVersion VersionNo) int {
	if vn.Major < anotherVersion.Major {
		return -1
	}
	if vn.Major > anotherVersion.Major {
		return 1
	}
	if vn.Minor < anotherVersion.Minor {
		return -1
	}
	if vn.Minor > anotherVersion.Minor {
		return 1
	}
	if vn.Revision < anotherVersion.Revision {
		return -1
	}
	if vn.Revision > anotherVersion.Revision {
		return 1
	}
	return 0
}

func (vn VersionNo) String() string {
	if vn.Revision == 0 {
		return fmt.Sprintf("%d.%d", vn.Major, vn.Minor)
	}
	return fmt.Sprintf("%d.%d.%d", vn.Major, vn.Minor, vn.Revision)
}
<|endoftext|>


<commit_before>package GoSDK

import "fmt"

const (
	_PORTALS_DEV_PREAMBLE  = "/admin/portals/"
	_PORTALS_USER_PREAMBLE = "/api/v/2/portals/"
)

func (d *DevClient) GetPortals(systemKey string) ([]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := get(d, _PORTALS_DEV_PREAMBLE+systemKey, nil, creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return nil, err
	}
	return resp.Body.([]interface{}), nil
}

func (d *DevClient) GetPortal(systemKey, name string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := get(d, _PORTALS_USER_PREAMBLE+systemKey+"/"+name, nil, creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) CreatePortal(systemKey, name string, dash map[string]interface{}) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}

	resp, err := post(d, _PORTALS_USER_PREAMBLE+systemKey+"/"+name, make(map[string]interface{}), creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return nil, err
	}
	// An empty portal has now been created.. populate with provided data
	return d.UpdatePortal(systemKey, name, dash)
}

func (d *DevClient) UpdatePortal(systemKey, name string, dash map[string]interface{}) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := put(d, _PORTALS_USER_PREAMBLE+systemKey+"/"+name, dash, creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) DeletePortal(systemKey, name string) error {
	creds, err := d.credentials()
	if err != nil {
		return err
	}
	resp, err := delete(d, _PORTALS_USER_PREAMBLE+systemKey+"/"+name, nil, creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>remove unused fmt import<commit_after>package GoSDK

const (
	_PORTALS_DEV_PREAMBLE  = "/admin/portals/"
	_PORTALS_USER_PREAMBLE = "/api/v/2/portals/"
)

func (d *DevClient) GetPortals(systemKey string) ([]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := get(d, _PORTALS_DEV_PREAMBLE+systemKey, nil, creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return nil, err
	}
	return resp.Body.([]interface{}), nil
}

func (d *DevClient) GetPortal(systemKey, name string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := get(d, _PORTALS_USER_PREAMBLE+systemKey+"/"+name, nil, creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) CreatePortal(systemKey, name string, dash map[string]interface{}) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}

	resp, err := post(d, _PORTALS_USER_PREAMBLE+systemKey+"/"+name, make(map[string]interface{}), creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return nil, err
	}
	// An empty portal has now been created.. populate with provided data
	return d.UpdatePortal(systemKey, name, dash)
}

func (d *DevClient) UpdatePortal(systemKey, name string, dash map[string]interface{}) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := put(d, _PORTALS_USER_PREAMBLE+systemKey+"/"+name, dash, creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) DeletePortal(systemKey, name string) error {
	creds, err := d.credentials()
	if err != nil {
		return err
	}
	resp, err := delete(d, _PORTALS_USER_PREAMBLE+systemKey+"/"+name, nil, creds, nil)
	resp, err = mapResponse(resp, err)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package posts

import (
	"time"
	"bufio"
	"os"
	"io/ioutil"
	"log"
	"sort"
	"encoding/json"
	"github.com/russross/blackfriday"
	"github.com/microcosm-cc/bluemonday"
)
type Post struct {
	Title, Subtitle string
	PublishDate time.Time
	MDContent string
	HTMLContent string
}
type ByAge []Post

func (b ByAge) Len() int {
	return len(b)
}
func (b ByAge) Swap(i, j int) {
	b[i], b[j] = b[j], b[i]
}
func (b ByAge) Less(i, j int) bool {
	return b[i].PublishDate.Unix() < b[j].PublishDate.Unix()
}

var PublicPosts []Post

func (p *Post) Render() {
	output := blackfriday.MarkdownCommon([]byte(p.MDContent))
	p.HTMLContent = string(bluemonday.UGCPolicy().SanitizeBytes([]byte(output)))
}

func NewPost(title, subtitle string, date time.Time, content string) Post {
	p := Post{title, subtitle, date, content, ""}
	p.Render()
	return p
}

func Load(folder string) {
	PublicPosts = make([]Post, 0)
	dirEntries, readError := ioutil.ReadDir(folder + "/posts")
	if readError != nil {
		log.Fatal(readError)
		os.Exit(1)
	}
	for _, entry := range dirEntries {
		if entry.IsDir() {
			continue
		}

		input, err := os.Open(folder + "/posts/" + entry.Name())
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
		defer input.Close()

		inputScanner := bufio.NewScanner(input)
		inputScanner.Split(bufio.ScanLines)

		header, body, headerActive := "{\n", "", true
		for inputScanner.Scan() {
			line := inputScanner.Text()
			log.Println("Read line", line)
			if headerActive {
				if line == "---" {
					headerActive = false
				} else {
					header += line + "\n"
				}
			} else {
				body += line + "\n";
			}
		}
		header += "}"
		log.Println(header)
		log.Println(body)

		var headerData struct {
			Title string `json:"title"`
			Subtitle string `json:"subtitle"`
			PublishDate string `json:"date"`
		}

		// load json header data
		jsonError := json.Unmarshal([]byte(header), &headerData)
		if jsonError != nil {
			log.Fatal(jsonError)
			os.Exit(1)
		}

		// parse Date from header data
		date, parseError := time.Parse("2015-12-31", headerData.PublishDate)
		if parseError != nil {
			log.Fatal(parseError)
			os.Exit(1)
		}
		PublicPosts = append(PublicPosts, NewPost(headerData.Title, headerData.Subtitle, date, body))
		log.Println("Read post file", entry.Name())
	}
	sort.Sort(ByAge(PublicPosts))
}

func GetNewest(count int) []Post {
	return PublicPosts[:count]
}
<commit_msg>Change time parsing<commit_after>package posts

import (
	"time"
	"bufio"
	"os"
	"io/ioutil"
	"log"
	"sort"
	"encoding/json"
	"github.com/russross/blackfriday"
	"github.com/microcosm-cc/bluemonday"
)
type Post struct {
	Title, Subtitle string
	PublishDate time.Time
	MDContent string
	HTMLContent string
}
type ByAge []Post

func (b ByAge) Len() int {
	return len(b)
}
func (b ByAge) Swap(i, j int) {
	b[i], b[j] = b[j], b[i]
}
func (b ByAge) Less(i, j int) bool {
	return b[i].PublishDate.Unix() < b[j].PublishDate.Unix()
}

var PublicPosts []Post

func (p *Post) Render() {
	output := blackfriday.MarkdownCommon([]byte(p.MDContent))
	p.HTMLContent = string(bluemonday.UGCPolicy().SanitizeBytes([]byte(output)))
}

func NewPost(title, subtitle string, date time.Time, content string) Post {
	p := Post{title, subtitle, date, content, ""}
	p.Render()
	return p
}

func Load(folder string) {
	PublicPosts = make([]Post, 0)
	dirEntries, readError := ioutil.ReadDir(folder + "/posts")
	if readError != nil {
		log.Fatal(readError)
		os.Exit(1)
	}
	for _, entry := range dirEntries {
		if entry.IsDir() {
			continue
		}

		input, err := os.Open(folder + "/posts/" + entry.Name())
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
		defer input.Close()

		inputScanner := bufio.NewScanner(input)
		inputScanner.Split(bufio.ScanLines)

		header, body, headerActive := "{\n", "", true
		for inputScanner.Scan() {
			line := inputScanner.Text()
			log.Println("Read line", line)
			if headerActive {
				if line == "---" {
					headerActive = false
				} else {
					header += line + "\n"
				}
			} else {
				body += line + "\n";
			}
		}
		header += "}"
		log.Println(header)
		log.Println(body)

		var headerData struct {
			Title string `json:"title"`
			Subtitle string `json:"subtitle"`
			PublishDate string `json:"date"`
		}

		// load json header data
		jsonError := json.Unmarshal([]byte(header), &headerData)
		if jsonError != nil {
			log.Fatal(jsonError)
			os.Exit(1)
		}

		// parse Date from header data
		date, parseError := time.Parse("2006-Jan-02", headerData.PublishDate)
		if parseError != nil {
			log.Fatal(parseError)
			os.Exit(1)
		}
		PublicPosts = append(PublicPosts, NewPost(headerData.Title, headerData.Subtitle, date, body))
		log.Println("Read post file", entry.Name())
	}
	sort.Sort(ByAge(PublicPosts))
}

func GetNewest(count int) []Post {
	return PublicPosts[:count]
}
<|endoftext|>


<commit_before>// Package analytics provides the methods to run an analytics reporting system
// for API requests which may be useful to users for measuring access and
// possibly identifying bad actors abusing requests.
package analytics

import (
	"encoding/json"
	"log"
	"net/http"
	"runtime"
	"strings"
	"time"

	"github.com/boltdb/bolt"
)

type apiRequest struct {
	URL        string `json:"url"`
	Method     string `json:"http_method"`
	Origin     string `json:"origin"`
	Proto      string `json:"http_protocol"`
	RemoteAddr string `json:"ip_address"`
	Timestamp  int64  `json:"timestamp"`
	External   bool   `json:"external"`
}

type apiMetric struct {
	Date   string `json:"date"`
	Total  int    `json:"total"`
	Unique int    `json:"unique"`
}

var (
	store       *bolt.DB
	requestChan chan apiRequest
)

// Record queues an apiRequest for metrics
func Record(req *http.Request) {
	external := strings.Contains(req.URL.Path, "/external/")

	r := apiRequest{
		URL:        req.URL.String(),
		Method:     req.Method,
		Origin:     req.Header.Get("Origin"),
		Proto:      req.Proto,
		RemoteAddr: req.RemoteAddr,
		Timestamp:  time.Now().Unix() * 1000,
		External:   external,
	}

	// put r on buffered requestChan to take advantage of batch insertion in DB
	requestChan <- r
}

// Close exports the abillity to close our db file. Should be called with defer
// after call to Init() from the same place.
func Close() {
	err := store.Close()
	if err != nil {
		log.Println(err)
	}
}

// Init creates a db connection, initializes the db with schema and data and
// sets up the queue/batching channel
func Init() {
	var err error
	store, err = bolt.Open("analytics.db", 0666, nil)
	if err != nil {
		log.Fatalln(err)
	}

	err = store.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("__requests"))
		if err != nil {
			return err
		}

		_, err = tx.CreateBucketIfNotExists([]byte("__metrics"))
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		log.Fatalln("Error idempotently creating requests bucket in analytics.db:", err)
	}

	requestChan = make(chan apiRequest, 1024*64*runtime.NumCPU())

	go serve()

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

func serve() {
	// make timer to notify select to batch request insert from requestChan
	// interval: 30 seconds
	apiRequestTimer := time.NewTicker(time.Second * 30)

	// make timer to notify select to remove analytics older than 14 days
	// interval: 1 week
	// TODO: enable analytics backup service to cloud
	pruneThreshold := time.Hour * 24 * 14
	pruneDBTimer := time.NewTicker(pruneThreshold / 2)

	for {
		select {
		case <-apiRequestTimer.C:
			err := batchInsert(requestChan)
			if err != nil {
				log.Println(err)
			}

		case <-pruneDBTimer.C:
			err := batchPrune(pruneThreshold)
			if err != nil {
				log.Println(err)
			}

		case <-time.After(time.Second * 30):
			continue
		}
	}
}

// ChartData returns the map containing decoded javascript needed to chart 2 weeks of data by day
func ChartData() (map[string]interface{}, error) {
	// set thresholds for today and the 13 days preceeding
	times := [14]time.Time{}
	dates := [14]string{}
	now := time.Now()
	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)

	ips := [14]map[string]struct{}{}
	for i := range ips {
		ips[i] = make(map[string]struct{})
	}

	total := [14]int{}
	unique := [14]int{}

	for i := range times {
		// subtract 24 * i hours to make days prior
		dur := time.Duration(24 * i * -1)
		day := today.Add(time.Hour * dur)

		// day threshold is [...n-1-i, n-1, n]
		times[len(times)-1-i] = day
		dates[len(times)-1-i] = day.Format("01/02")
	}

	// get api request analytics and metrics from db
	var requests = []apiRequest{}
	var metrics = [14]apiMetric{}

	err := store.View(func(tx *bolt.Tx) error {
		m := tx.Bucket([]byte("__metrics"))
		b := tx.Bucket([]byte("__requests"))

		err := m.ForEach(func(k, v []byte) error {
			var metric apiMetric
			err := json.Unmarshal(v, &metric)
			if err != nil {
				log.Println("Error decoding api metric json from analytics db:", err)
				return nil
			}

			// if the metric date is in current date range, insert it into
			// metrics array at the position of the date in dates array
			for i := range dates {
				if metric.Date == dates[i] {
					metrics[i] = metric
				}
			}

			return nil
		})
		if err != nil {
			return err
		}

		err = b.ForEach(func(k, v []byte) error {
			var r apiRequest
			err := json.Unmarshal(v, &r)
			if err != nil {
				log.Println("Error decoding api request json from analytics db:", err)
				return nil
			}

			// delete the record in db if it belongs to a day already in metrics,
			// otherwise append it to requests to be analyzed
			d := time.Unix(r.Timestamp/1000, 0).Format("01/02")
			for i := range dates {
				if metrics[i].Date == d {
					err := b.Delete(k)
					if err != nil {
						return err
					}
				} else {
					requests = append(requests, r)
				}
			}

			return nil
		})
		if err != nil {
			return err
		}

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

CHECK_REQUEST:
	for i := range requests {
		ts := time.Unix(requests[i].Timestamp/1000, 0)

		for j := range times {
			// if on today, there will be no next iteration to set values for
			// day prior so all valid requests belong to today
			if j == len(times)-1 {
				if ts.After(times[j]) || ts.Equal(times[j]) {
					// do all record keeping
					total[j]++

					if _, ok := ips[j][requests[i].RemoteAddr]; !ok {
						unique[j]++
						ips[j][requests[i].RemoteAddr] = struct{}{}
					}

					continue CHECK_REQUEST
				}
			}

			if ts.Equal(times[j]) {
				// increment total count for current time threshold (day)
				total[j]++

				// if no IP found for current threshold, increment unique and record IP
				if _, ok := ips[j][requests[i].RemoteAddr]; !ok {
					unique[j]++
					ips[j][requests[i].RemoteAddr] = struct{}{}
				}

				continue CHECK_REQUEST
			}

			if ts.Before(times[j]) {
				// check if older than earliest threshold
				if j == 0 {
					continue CHECK_REQUEST
				}

				// increment total count for previous time threshold (day)
				total[j-1]++

				// if no IP found for day prior, increment unique and record IP
				if _, ok := ips[j-1][requests[i].RemoteAddr]; !ok {
					unique[j-1]++
					ips[j-1][requests[i].RemoteAddr] = struct{}{}
				}
			}
		}
	}

	// loop through total and unique to see which dates are accounted for and
	// insert data from metrics array where dates are not
	for i := range metrics {
		if total[i] == 0 {
			total[i] = metrics[i].Total
		}

		if unique[i] == 0 {
			unique[i] = metrics[i].Unique
		}
	}

	// marshal array counts to js arrays for output to chart
	jsUnique, err := json.Marshal(unique)
	if err != nil {
		return nil, err
	}

	jsTotal, err := json.Marshal(total)
	if err != nil {
		return nil, err
	}

	return map[string]interface{}{
		"dates":  dates,
		"unique": string(jsUnique),
		"total":  string(jsTotal),
		"from":   dates[0],
		"to":     dates[len(dates)-1],
	}, nil
}
<commit_msg>checking for metric data in db first<commit_after>// Package analytics provides the methods to run an analytics reporting system
// for API requests which may be useful to users for measuring access and
// possibly identifying bad actors abusing requests.
package analytics

import (
	"encoding/json"
	"log"
	"net/http"
	"runtime"
	"strings"
	"time"

	"github.com/boltdb/bolt"
)

type apiRequest struct {
	URL        string `json:"url"`
	Method     string `json:"http_method"`
	Origin     string `json:"origin"`
	Proto      string `json:"http_protocol"`
	RemoteAddr string `json:"ip_address"`
	Timestamp  int64  `json:"timestamp"`
	External   bool   `json:"external"`
}

type apiMetric struct {
	Date   string `json:"date"`
	Total  int    `json:"total"`
	Unique int    `json:"unique"`
}

var (
	store       *bolt.DB
	requestChan chan apiRequest
)

// Record queues an apiRequest for metrics
func Record(req *http.Request) {
	external := strings.Contains(req.URL.Path, "/external/")

	r := apiRequest{
		URL:        req.URL.String(),
		Method:     req.Method,
		Origin:     req.Header.Get("Origin"),
		Proto:      req.Proto,
		RemoteAddr: req.RemoteAddr,
		Timestamp:  time.Now().Unix() * 1000,
		External:   external,
	}

	// put r on buffered requestChan to take advantage of batch insertion in DB
	requestChan <- r
}

// Close exports the abillity to close our db file. Should be called with defer
// after call to Init() from the same place.
func Close() {
	err := store.Close()
	if err != nil {
		log.Println(err)
	}
}

// Init creates a db connection, initializes the db with schema and data and
// sets up the queue/batching channel
func Init() {
	var err error
	store, err = bolt.Open("analytics.db", 0666, nil)
	if err != nil {
		log.Fatalln(err)
	}

	err = store.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("__requests"))
		if err != nil {
			return err
		}

		_, err = tx.CreateBucketIfNotExists([]byte("__metrics"))
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		log.Fatalln("Error idempotently creating requests bucket in analytics.db:", err)
	}

	requestChan = make(chan apiRequest, 1024*64*runtime.NumCPU())

	go serve()

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

func serve() {
	// make timer to notify select to batch request insert from requestChan
	// interval: 30 seconds
	apiRequestTimer := time.NewTicker(time.Second * 30)

	// make timer to notify select to remove analytics older than 14 days
	// interval: 1 week
	// TODO: enable analytics backup service to cloud
	pruneThreshold := time.Hour * 24 * 14
	pruneDBTimer := time.NewTicker(pruneThreshold / 2)

	for {
		select {
		case <-apiRequestTimer.C:
			err := batchInsert(requestChan)
			if err != nil {
				log.Println(err)
			}

		case <-pruneDBTimer.C:
			err := batchPrune(pruneThreshold)
			if err != nil {
				log.Println(err)
			}

		case <-time.After(time.Second * 30):
			continue
		}
	}
}

// ChartData returns the map containing decoded javascript needed to chart 2 weeks of data by day
func ChartData() (map[string]interface{}, error) {
	// set thresholds for today and the 13 days preceeding
	times := [14]time.Time{}
	dates := [14]string{}
	now := time.Now()
	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)

	ips := [14]map[string]struct{}{}
	for i := range ips {
		ips[i] = make(map[string]struct{})
	}

	total := [14]int{}
	unique := [14]int{}

	for i := range times {
		// subtract 24 * i hours to make days prior
		dur := time.Duration(24 * i * -1)
		day := today.Add(time.Hour * dur)

		// day threshold is [...n-1-i, n-1, n]
		times[len(times)-1-i] = day
		dates[len(times)-1-i] = day.Format("01/02")
	}

	// get api request analytics and metrics from db
	var requests = []apiRequest{}
	var metrics = [14]apiMetric{}

	err := store.View(func(tx *bolt.Tx) error {
		m := tx.Bucket([]byte("__metrics"))
		b := tx.Bucket([]byte("__requests"))

		err := m.ForEach(func(k, v []byte) error {
			var metric apiMetric
			err := json.Unmarshal(v, &metric)
			if err != nil {
				log.Println("Error decoding api metric json from analytics db:", err)
				return nil
			}

			// if the metric date is in current date range, insert it into
			// metrics array at the position of the date in dates array
			for i := range dates {
				if metric.Date == dates[i] {
					metrics[i] = metric
				}
			}

			return nil
		})
		if err != nil {
			return err
		}

		err = b.ForEach(func(k, v []byte) error {
			var r apiRequest
			err := json.Unmarshal(v, &r)
			if err != nil {
				log.Println("Error decoding api request json from analytics db:", err)
				return nil
			}

			// delete the record in db if it belongs to a day already in metrics,
			// otherwise append it to requests to be analyzed
			d := time.Unix(r.Timestamp/1000, 0).Format("01/02")
			if m.Get([]byte(d)) != nil {
				err := b.Delete(k)
				if err != nil {
					return err
				}
			} else {
				requests = append(requests, r)
			}

			return nil
		})
		if err != nil {
			return err
		}

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

CHECK_REQUEST:
	for i := range requests {
		ts := time.Unix(requests[i].Timestamp/1000, 0)

		for j := range times {
			// if on today, there will be no next iteration to set values for
			// day prior so all valid requests belong to today
			if j == len(times)-1 {
				if ts.After(times[j]) || ts.Equal(times[j]) {
					// do all record keeping
					total[j]++

					if _, ok := ips[j][requests[i].RemoteAddr]; !ok {
						unique[j]++
						ips[j][requests[i].RemoteAddr] = struct{}{}
					}

					continue CHECK_REQUEST
				}
			}

			if ts.Equal(times[j]) {
				// increment total count for current time threshold (day)
				total[j]++

				// if no IP found for current threshold, increment unique and record IP
				if _, ok := ips[j][requests[i].RemoteAddr]; !ok {
					unique[j]++
					ips[j][requests[i].RemoteAddr] = struct{}{}
				}

				continue CHECK_REQUEST
			}

			if ts.Before(times[j]) {
				// check if older than earliest threshold
				if j == 0 {
					continue CHECK_REQUEST
				}

				// increment total count for previous time threshold (day)
				total[j-1]++

				// if no IP found for day prior, increment unique and record IP
				if _, ok := ips[j-1][requests[i].RemoteAddr]; !ok {
					unique[j-1]++
					ips[j-1][requests[i].RemoteAddr] = struct{}{}
				}
			}
		}
	}

	// loop through total and unique to see which dates are accounted for and
	// insert data from metrics array where dates are not
	for i := range metrics {
		if total[i] == 0 {
			total[i] = metrics[i].Total
		}

		if unique[i] == 0 {
			unique[i] = metrics[i].Unique
		}
	}

	// marshal array counts to js arrays for output to chart
	jsUnique, err := json.Marshal(unique)
	if err != nil {
		return nil, err
	}

	jsTotal, err := json.Marshal(total)
	if err != nil {
		return nil, err
	}

	return map[string]interface{}{
		"dates":  dates,
		"unique": string(jsUnique),
		"total":  string(jsTotal),
		"from":   dates[0],
		"to":     dates[len(dates)-1],
	}, nil
}
<|endoftext|>


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

package proto

import (
	"code.google.com/p/snappy-go/snappy"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"hash"
	"io"
	"sync"
)

type CommandIO struct {
	writeAuth   hash.Hash
	cryptWriter io.Writer
	readAuth    hash.Hash
	cryptReader io.Reader
	conn        io.ReadWriter

	writeLock *sync.Mutex
}

func (self *CommandIO) writeThenHmac(data []byte, encrypt bool) (mac []byte, err error) {
	writer := self.cryptWriter
	if !encrypt {
		writer = self.conn
	} else {
		self.writeAuth.Reset()
	}
	err = writen(writer, data)
	if err != nil {
		return
	}
	if !encrypt {
		return
	}
	mac = self.writeAuth.Sum(nil)
	return
}

func (self *CommandIO) readThenHmac(data []byte, encrypt bool) (mac []byte, err error) {
	reader := self.cryptReader
	if !encrypt {
		reader = self.conn
	} else {
		self.readAuth.Reset()
	}

	n, err := io.ReadFull(reader, data)
	if err != nil {
		return
	}
	if n != len(data) {
		err = io.EOF
		return
	}
	if !encrypt {
		return
	}
	mac = self.readAuth.Sum(nil)
	return
}

func (self *CommandIO) writeHmac(mac []byte) error {
	if len(mac) == 0 {
		return nil
	}
	return writen(self.conn, mac)
}

func (self *CommandIO) readAndCmpHmac(mac []byte) error {
	if len(mac) == 0 {
		return nil
	}
	macRecved := make([]byte, self.readAuth.Size())
	n, err := io.ReadFull(self.conn, macRecved)
	if err != nil {
		return err
	}
	if n != len(macRecved) {
		return ErrCorruptedData
	}
	if !bytesEq(mac, macRecved) {
		return ErrCorruptedData
	}
	return nil
}

func (self *CommandIO) decodeCommand(data []byte, compress bool) (cmd *Command, err error) {
	decoded := data
	if !compress {
		decoded, err = snappy.Decode(nil, data)
		if err != nil {
			return
		}
	}
	cmd, err = UnmarshalCommand(decoded)
	if err != nil {
		return
	}
	return
}

func (self *CommandIO) encodeCommand(cmd *Command, compress bool) (data []byte, err error) {
	bsonEncoded, err := cmd.Marshal()
	if err != nil {
		return
	}

	data = bsonEncoded
	if !compress {
		data, err = snappy.Encode(nil, bsonEncoded)
		if err != nil {
			return
		}
	}
	return
}

// WriteCommand() is goroutine-safe. i.e. Multiple goroutine could write concurrently.
func (self *CommandIO) WriteCommand(cmd *Command, compress, encrypt bool) error {
	var flag uint16
	flag = 0
	if compress {
		flag |= cmdflag_COMPRESS
	}
	if encrypt {
		flag |= cmdflag_ENCRYPT
	}
	data, err := self.encodeCommand(cmd, compress)
	if err != nil {
		return err
	}
	var cmdLen uint16
	cmdLen = uint16(len(data))
	if cmdLen == 0 {
		return nil
	}
	self.writeLock.Lock()
	defer self.writeLock.Unlock()
	err = binary.Write(self.conn, binary.LittleEndian, cmdLen)
	if err != nil {
		return err
	}
	err = binary.Write(self.conn, binary.LittleEndian, flag)
	if err != nil {
		return err
	}
	mac, err := self.writeThenHmac(data, encrypt)
	if err != nil {
		return err
	}
	err = self.writeHmac(mac)
	if err != nil {
		return err
	}
	return nil
}

// ReadCommand() is not goroutine-safe.
func (self *CommandIO) ReadCommand() (cmd *Command, err error) {
	var cmdLen uint16
	var flag uint16
	err = binary.Read(self.conn, binary.LittleEndian, &cmdLen)
	if err != nil {
		return
	}
	err = binary.Read(self.conn, binary.LittleEndian, &flag)
	if err != nil {
		return
	}

	compress := ((flag & cmdflag_COMPRESS) != 0)
	encrypt := ((flag & cmdflag_ENCRYPT) != 0)

	data := make([]byte, int(cmdLen))
	mac, err := self.readThenHmac(data, encrypt)
	if err != nil {
		return
	}
	err = self.readAndCmpHmac(mac)
	if err != nil {
		return
	}
	cmd, err = self.decodeCommand(data, compress)
	return
}

func NewCommandIO(writeKey, writeAuthKey, readKey, readAuthKey []byte, conn io.ReadWriter) *CommandIO {
	ret := new(CommandIO)
	ret.writeAuth = hmac.New(sha256.New, writeAuthKey)
	ret.readAuth = hmac.New(sha256.New, readAuthKey)
	ret.conn = conn
	ret.writeLock = new(sync.Mutex)

	writeBlkCipher, _ := aes.NewCipher(writeKey)
	readBlkCipher, _ := aes.NewCipher(readKey)

	// IV: 0 for all. Since we change keys for each connection, letting IV=0 won't hurt.
	writeIV := make([]byte, writeBlkCipher.BlockSize())
	readIV := make([]byte, readBlkCipher.BlockSize())

	writeStream := cipher.NewCTR(writeBlkCipher, writeIV)
	readStream := cipher.NewCTR(readBlkCipher, readIV)

	// Then for each encrypted bit,
	// it will be written to both the connection and the hmac
	// We use encrypt-then-hmac scheme.
	mwriter := io.MultiWriter(conn, ret.writeAuth)
	swriter := new(cipher.StreamWriter)
	swriter.S = writeStream
	swriter.W = mwriter
	ret.cryptWriter = swriter

	// Similarly, for each bit read from the connection,
	// it will be written to the hmac as well.
	tee := io.TeeReader(conn, ret.readAuth)
	sreader := new(cipher.StreamReader)
	sreader.S = readStream
	sreader.R = tee
	ret.cryptReader = sreader
	return ret
}
<commit_msg>compress flag was set to its opposite.<commit_after>/*
 * Copyright 2013 Nan Deng
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package proto

import (
	"code.google.com/p/snappy-go/snappy"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"hash"
	"io"
	"sync"
)

type CommandIO struct {
	writeAuth   hash.Hash
	cryptWriter io.Writer
	readAuth    hash.Hash
	cryptReader io.Reader
	conn        io.ReadWriter

	writeLock *sync.Mutex
}

func (self *CommandIO) writeThenHmac(data []byte, encrypt bool) (mac []byte, err error) {
	writer := self.cryptWriter
	if !encrypt {
		writer = self.conn
	} else {
		self.writeAuth.Reset()
	}
	err = writen(writer, data)
	if err != nil {
		return
	}
	if !encrypt {
		return
	}
	mac = self.writeAuth.Sum(nil)
	return
}

func (self *CommandIO) readThenHmac(data []byte, encrypt bool) (mac []byte, err error) {
	reader := self.cryptReader
	if !encrypt {
		reader = self.conn
	} else {
		self.readAuth.Reset()
	}

	n, err := io.ReadFull(reader, data)
	if err != nil {
		return
	}
	if n != len(data) {
		err = io.EOF
		return
	}
	if !encrypt {
		return
	}
	mac = self.readAuth.Sum(nil)
	return
}

func (self *CommandIO) writeHmac(mac []byte) error {
	if len(mac) == 0 {
		return nil
	}
	return writen(self.conn, mac)
}

func (self *CommandIO) readAndCmpHmac(mac []byte) error {
	if len(mac) == 0 {
		return nil
	}
	macRecved := make([]byte, self.readAuth.Size())
	n, err := io.ReadFull(self.conn, macRecved)
	if err != nil {
		return err
	}
	if n != len(macRecved) {
		return ErrCorruptedData
	}
	if !bytesEq(mac, macRecved) {
		return ErrCorruptedData
	}
	return nil
}

func (self *CommandIO) decodeCommand(data []byte, compress bool) (cmd *Command, err error) {
	decoded := data
	if compress {
		decoded, err = snappy.Decode(nil, data)
		if err != nil {
			return
		}
	}
	cmd, err = UnmarshalCommand(decoded)
	if err != nil {
		return
	}
	return
}

func (self *CommandIO) encodeCommand(cmd *Command, compress bool) (data []byte, err error) {
	bsonEncoded, err := cmd.Marshal()
	if err != nil {
		return
	}

	data = bsonEncoded
	if compress {
		data, err = snappy.Encode(nil, bsonEncoded)
		if err != nil {
			return
		}
	}
	return
}

// WriteCommand() is goroutine-safe. i.e. Multiple goroutine could write concurrently.
func (self *CommandIO) WriteCommand(cmd *Command, compress, encrypt bool) error {
	var flag uint16
	flag = 0
	if compress {
		flag |= cmdflag_COMPRESS
	}
	if encrypt {
		flag |= cmdflag_ENCRYPT
	}
	data, err := self.encodeCommand(cmd, compress)
	if err != nil {
		return err
	}
	var cmdLen uint16
	cmdLen = uint16(len(data))
	if cmdLen == 0 {
		return nil
	}
	self.writeLock.Lock()
	defer self.writeLock.Unlock()
	err = binary.Write(self.conn, binary.LittleEndian, cmdLen)
	if err != nil {
		return err
	}
	err = binary.Write(self.conn, binary.LittleEndian, flag)
	if err != nil {
		return err
	}
	mac, err := self.writeThenHmac(data, encrypt)
	if err != nil {
		return err
	}
	err = self.writeHmac(mac)
	if err != nil {
		return err
	}
	return nil
}

// ReadCommand() is not goroutine-safe.
func (self *CommandIO) ReadCommand() (cmd *Command, err error) {
	var cmdLen uint16
	var flag uint16
	err = binary.Read(self.conn, binary.LittleEndian, &cmdLen)
	if err != nil {
		return
	}
	err = binary.Read(self.conn, binary.LittleEndian, &flag)
	if err != nil {
		return
	}

	compress := ((flag & cmdflag_COMPRESS) != 0)
	encrypt := ((flag & cmdflag_ENCRYPT) != 0)

	data := make([]byte, int(cmdLen))
	mac, err := self.readThenHmac(data, encrypt)
	if err != nil {
		return
	}
	err = self.readAndCmpHmac(mac)
	if err != nil {
		return
	}
	cmd, err = self.decodeCommand(data, compress)
	return
}

func NewCommandIO(writeKey, writeAuthKey, readKey, readAuthKey []byte, conn io.ReadWriter) *CommandIO {
	ret := new(CommandIO)
	ret.writeAuth = hmac.New(sha256.New, writeAuthKey)
	ret.readAuth = hmac.New(sha256.New, readAuthKey)
	ret.conn = conn
	ret.writeLock = new(sync.Mutex)

	writeBlkCipher, _ := aes.NewCipher(writeKey)
	readBlkCipher, _ := aes.NewCipher(readKey)

	// IV: 0 for all. Since we change keys for each connection, letting IV=0 won't hurt.
	writeIV := make([]byte, writeBlkCipher.BlockSize())
	readIV := make([]byte, readBlkCipher.BlockSize())

	writeStream := cipher.NewCTR(writeBlkCipher, writeIV)
	readStream := cipher.NewCTR(readBlkCipher, readIV)

	// Then for each encrypted bit,
	// it will be written to both the connection and the hmac
	// We use encrypt-then-hmac scheme.
	mwriter := io.MultiWriter(conn, ret.writeAuth)
	swriter := new(cipher.StreamWriter)
	swriter.S = writeStream
	swriter.W = mwriter
	ret.cryptWriter = swriter

	// Similarly, for each bit read from the connection,
	// it will be written to the hmac as well.
	tee := io.TeeReader(conn, ret.readAuth)
	sreader := new(cipher.StreamReader)
	sreader.S = readStream
	sreader.R = tee
	ret.cryptReader = sreader
	return ret
}
<|endoftext|>


<commit_before>package system

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	httpd "github.com/rqlite/rqlite/http"
	"github.com/rqlite/rqlite/store"
)

// Node represents a node under test.
type Node struct {
	APIAddr  string
	RaftAddr string
	Dir      string
	Store    *store.Store
	Service  *httpd.Service
}

// SameAs returns true if this node is the same as o node.
func (n *Node) SameAs(o *Node) bool {
	return n.RaftAddr == o.RaftAddr
}

// Deprovision shuts down and removes all resources associated with the node.
func (n *Node) Deprovision() {
	n.Store.Close(false)
	n.Service.Close()
	os.RemoveAll(n.Dir)
}

// WaitForLeader blocks for up to 10 seconds until the node detects a leader.
func (n *Node) WaitForLeader() (string, error) {
	return n.Store.WaitForLeader(10 * time.Second)
}

// Execute executes a single statement against the node.
func (n *Node) Execute(stmt string) (string, error) {
	return n.ExecuteMulti([]string{stmt})
}

// ExecuteMulti executes multiple statements against the node.
func (n *Node) ExecuteMulti(stmts []string) (string, error) {
	j, err := json.Marshal(stmts)
	if err != nil {
		return "", err
	}
	return n.postExecute(string(j))
}

// Query runs a single query against the node.
func (n *Node) Query(stmt string) (string, error) {
	v, _ := url.Parse("http://" + n.APIAddr + "/db/query")
	v.RawQuery = url.Values{"q": []string{stmt}}.Encode()

	resp, err := http.Get(v.String())
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	return string(body), nil
}

// Join instructs this node to join the leader.
func (n *Node) Join(leader *Node) error {
	b, err := json.Marshal(map[string]string{"addr": n.RaftAddr})
	if err != nil {
		return err
	}

	// Attempt to join leader
	resp, err := http.Post("http://"+leader.APIAddr+"/join", "application-type/json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("failed to join, leader returned: %s", resp.Status)
	}
	defer resp.Body.Close()
	return nil
}

// ConfirmRedirect confirms that the node responds with a redirect to the given host.
func (n *Node) ConfirmRedirect(host string) bool {
	v, _ := url.Parse("http://" + n.APIAddr + "/db/query")
	v.RawQuery = url.Values{"q": []string{`SELECT * FROM foo`}}.Encode()

	resp, err := http.Get(v.String())
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	fmt.Println(resp.StatusCode)
	if resp.StatusCode != http.StatusMovedPermanently {
		return false
	}
	if resp.Header.Get("location") != host {
		return false
	}
	return true
}

func (n *Node) postExecute(stmt string) (string, error) {
	resp, err := http.Post("http://"+n.APIAddr+"/db/execute", "application/json", strings.NewReader(stmt))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	return string(body), nil
}

// Cluster represents a cluster of nodes.
type Cluster []*Node

// Leader returns the leader node of a cluster.
func (c Cluster) Leader() (*Node, error) {
	l, err := c[0].WaitForLeader()
	if err != nil {
		return nil, err
	}
	return c.FindNodeByRaftAddr(l)
}

// WaitForNewLeader waits for the leader to change from the node passed in.
func (c Cluster) WaitForNewLeader(old *Node) (*Node, error) {
	timer := time.NewTimer(30 * time.Second)
	defer timer.Stop()
	ticker := time.NewTicker(100 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-timer.C:
			return nil, fmt.Errorf("timed out waiting for new leader")
		case <-ticker.C:
			l, err := c.Leader()
			if err != nil {
				continue
			}
			if !l.SameAs(old) {
				return l, nil
			}
		}
	}
}

// Followers returns the slide of nodes in the cluster that are followers.
func (c Cluster) Followers() ([]*Node, error) {
	n, err := c[0].WaitForLeader()
	if err != nil {
		return nil, err
	}
	leader, err := c.FindNodeByRaftAddr(n)
	if err != nil {
		return nil, err
	}

	var followers []*Node
	for _, n := range c {
		if n != leader {
			followers = append(followers, n)
		}
	}
	return followers, nil
}

// RemoveNode removes the given node from the list of nodes representing
// a cluster.
func (c Cluster) RemoveNode(node *Node) {
	for i, n := range c {
		if n.RaftAddr == node.RaftAddr {
			c = append(c[:i], c[i+1:]...)
			return
		}
	}
}

// FindNodeByRaftAddr returns the node with the given Raft address.
func (c Cluster) FindNodeByRaftAddr(addr string) (*Node, error) {
	for _, n := range c {
		if n.RaftAddr == addr {
			return n, nil
		}
	}
	return nil, fmt.Errorf("node not found")
}

// Deprovision deprovisions every node in the cluster.
func (c Cluster) Deprovision() {
	for _, n := range c {
		n.Deprovision()
	}
}

// Remove tells the cluster at n to remove the node at addr. Assumes n is the leader.
func Remove(n *Node, addr string) error {
	b, err := json.Marshal(map[string]string{"addr": addr})
	if err != nil {
		return err
	}

	// Attempt to remove node from leader.
	resp, err := http.Post("http://"+n.APIAddr+"/remove", "application-type/json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("failed to remove node, leader returned: %s", resp.Status)
	}
	defer resp.Body.Close()
	return nil
}

func mustNewNode(enableSingle bool) *Node {
	node := &Node{
		Dir: mustTempDir(),
	}

	dbConf := store.NewDBConfig("", false)
	node.Store = store.New(dbConf, node.Dir, mustMockTransport("localhost:0"))
	if err := node.Store.Open(enableSingle); err != nil {
		node.Deprovision()
		panic(fmt.Sprintf("failed to open store: %s", err.Error()))
	}
	node.RaftAddr = node.Store.Addr().String()

	node.Service = httpd.New("localhost:0", node.Store, nil)
	if err := node.Service.Start(); err != nil {
		node.Deprovision()
		panic(fmt.Sprintf("failed to start HTTP server: %s", err.Error()))
	}
	node.APIAddr = node.Service.Addr().String()

	return node
}

func mustNewLeaderNode() *Node {
	node := mustNewNode(true)
	if _, err := node.WaitForLeader(); err != nil {
		node.Deprovision()
		panic("node never became leader")
	}
	return node
}

type mockTransport struct {
	ln net.Listener
}

func mustMockTransport(addr string) *mockTransport {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		panic("failed to create new transport")
	}
	return &mockTransport{ln}
}

func (m *mockTransport) Dial(addr string, timeout time.Duration) (net.Conn, error) {
	return net.DialTimeout("tcp", addr, timeout)
}

func (m *mockTransport) Accept() (net.Conn, error) { return m.ln.Accept() }

func (m *mockTransport) Close() error { return m.ln.Close() }

func (m *mockTransport) Addr() net.Addr { return m.ln.Addr() }

func mustTempDir() string {
	var err error
	path, err := ioutil.TempDir("", "rqlilte-system-test-")
	if err != nil {
		panic("failed to create temp dir")
	}
	return path
}
<commit_msg>Typo in comment<commit_after>package system

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	httpd "github.com/rqlite/rqlite/http"
	"github.com/rqlite/rqlite/store"
)

// Node represents a node under test.
type Node struct {
	APIAddr  string
	RaftAddr string
	Dir      string
	Store    *store.Store
	Service  *httpd.Service
}

// SameAs returns true if this node is the same as node o.
func (n *Node) SameAs(o *Node) bool {
	return n.RaftAddr == o.RaftAddr
}

// Deprovision shuts down and removes all resources associated with the node.
func (n *Node) Deprovision() {
	n.Store.Close(false)
	n.Service.Close()
	os.RemoveAll(n.Dir)
}

// WaitForLeader blocks for up to 10 seconds until the node detects a leader.
func (n *Node) WaitForLeader() (string, error) {
	return n.Store.WaitForLeader(10 * time.Second)
}

// Execute executes a single statement against the node.
func (n *Node) Execute(stmt string) (string, error) {
	return n.ExecuteMulti([]string{stmt})
}

// ExecuteMulti executes multiple statements against the node.
func (n *Node) ExecuteMulti(stmts []string) (string, error) {
	j, err := json.Marshal(stmts)
	if err != nil {
		return "", err
	}
	return n.postExecute(string(j))
}

// Query runs a single query against the node.
func (n *Node) Query(stmt string) (string, error) {
	v, _ := url.Parse("http://" + n.APIAddr + "/db/query")
	v.RawQuery = url.Values{"q": []string{stmt}}.Encode()

	resp, err := http.Get(v.String())
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	return string(body), nil
}

// Join instructs this node to join the leader.
func (n *Node) Join(leader *Node) error {
	b, err := json.Marshal(map[string]string{"addr": n.RaftAddr})
	if err != nil {
		return err
	}

	// Attempt to join leader
	resp, err := http.Post("http://"+leader.APIAddr+"/join", "application-type/json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("failed to join, leader returned: %s", resp.Status)
	}
	defer resp.Body.Close()
	return nil
}

// ConfirmRedirect confirms that the node responds with a redirect to the given host.
func (n *Node) ConfirmRedirect(host string) bool {
	v, _ := url.Parse("http://" + n.APIAddr + "/db/query")
	v.RawQuery = url.Values{"q": []string{`SELECT * FROM foo`}}.Encode()

	resp, err := http.Get(v.String())
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	fmt.Println(resp.StatusCode)
	if resp.StatusCode != http.StatusMovedPermanently {
		return false
	}
	if resp.Header.Get("location") != host {
		return false
	}
	return true
}

func (n *Node) postExecute(stmt string) (string, error) {
	resp, err := http.Post("http://"+n.APIAddr+"/db/execute", "application/json", strings.NewReader(stmt))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	return string(body), nil
}

// Cluster represents a cluster of nodes.
type Cluster []*Node

// Leader returns the leader node of a cluster.
func (c Cluster) Leader() (*Node, error) {
	l, err := c[0].WaitForLeader()
	if err != nil {
		return nil, err
	}
	return c.FindNodeByRaftAddr(l)
}

// WaitForNewLeader waits for the leader to change from the node passed in.
func (c Cluster) WaitForNewLeader(old *Node) (*Node, error) {
	timer := time.NewTimer(30 * time.Second)
	defer timer.Stop()
	ticker := time.NewTicker(100 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-timer.C:
			return nil, fmt.Errorf("timed out waiting for new leader")
		case <-ticker.C:
			l, err := c.Leader()
			if err != nil {
				continue
			}
			if !l.SameAs(old) {
				return l, nil
			}
		}
	}
}

// Followers returns the slide of nodes in the cluster that are followers.
func (c Cluster) Followers() ([]*Node, error) {
	n, err := c[0].WaitForLeader()
	if err != nil {
		return nil, err
	}
	leader, err := c.FindNodeByRaftAddr(n)
	if err != nil {
		return nil, err
	}

	var followers []*Node
	for _, n := range c {
		if n != leader {
			followers = append(followers, n)
		}
	}
	return followers, nil
}

// RemoveNode removes the given node from the list of nodes representing
// a cluster.
func (c Cluster) RemoveNode(node *Node) {
	for i, n := range c {
		if n.RaftAddr == node.RaftAddr {
			c = append(c[:i], c[i+1:]...)
			return
		}
	}
}

// FindNodeByRaftAddr returns the node with the given Raft address.
func (c Cluster) FindNodeByRaftAddr(addr string) (*Node, error) {
	for _, n := range c {
		if n.RaftAddr == addr {
			return n, nil
		}
	}
	return nil, fmt.Errorf("node not found")
}

// Deprovision deprovisions every node in the cluster.
func (c Cluster) Deprovision() {
	for _, n := range c {
		n.Deprovision()
	}
}

// Remove tells the cluster at n to remove the node at addr. Assumes n is the leader.
func Remove(n *Node, addr string) error {
	b, err := json.Marshal(map[string]string{"addr": addr})
	if err != nil {
		return err
	}

	// Attempt to remove node from leader.
	resp, err := http.Post("http://"+n.APIAddr+"/remove", "application-type/json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("failed to remove node, leader returned: %s", resp.Status)
	}
	defer resp.Body.Close()
	return nil
}

func mustNewNode(enableSingle bool) *Node {
	node := &Node{
		Dir: mustTempDir(),
	}

	dbConf := store.NewDBConfig("", false)
	node.Store = store.New(dbConf, node.Dir, mustMockTransport("localhost:0"))
	if err := node.Store.Open(enableSingle); err != nil {
		node.Deprovision()
		panic(fmt.Sprintf("failed to open store: %s", err.Error()))
	}
	node.RaftAddr = node.Store.Addr().String()

	node.Service = httpd.New("localhost:0", node.Store, nil)
	if err := node.Service.Start(); err != nil {
		node.Deprovision()
		panic(fmt.Sprintf("failed to start HTTP server: %s", err.Error()))
	}
	node.APIAddr = node.Service.Addr().String()

	return node
}

func mustNewLeaderNode() *Node {
	node := mustNewNode(true)
	if _, err := node.WaitForLeader(); err != nil {
		node.Deprovision()
		panic("node never became leader")
	}
	return node
}

type mockTransport struct {
	ln net.Listener
}

func mustMockTransport(addr string) *mockTransport {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		panic("failed to create new transport")
	}
	return &mockTransport{ln}
}

func (m *mockTransport) Dial(addr string, timeout time.Duration) (net.Conn, error) {
	return net.DialTimeout("tcp", addr, timeout)
}

func (m *mockTransport) Accept() (net.Conn, error) { return m.ln.Accept() }

func (m *mockTransport) Close() error { return m.ln.Close() }

func (m *mockTransport) Addr() net.Addr { return m.ln.Addr() }

func mustTempDir() string {
	var err error
	path, err := ioutil.TempDir("", "rqlilte-system-test-")
	if err != nil {
		panic("failed to create temp dir")
	}
	return path
}
<|endoftext|>


<commit_before>package dynmap

import (
	"strings"
	// "log"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"time"
)

//Dont make this a map type, since we want the option of
//extending this and adding members.
type DynMap struct {
	Map map[string]interface{}
}

type DynMaper interface {
	ToDynMap() *DynMap
}

func NewDynMap() *DynMap {
	return &DynMap{make(map[string]interface{})}
}

// Recursively converts this to a regular go map.
// (will convert any sub maps)
func (this *DynMap) ToMap() map[string]interface{} {
	mp := make(map[string]interface{})
	for k, v := range(this.Map) {
		submp, ok := ToDynMap(this.Map[k])
		if ok {
			v = submp.ToMap()
		}
		mp[k] = v
	}
	return mp
}

// Returns self. Here so that we satisfy the DynMaper interface
func (this *DynMap) ToDynMap() *DynMap {
	return this
}

//encodes this map into a url encoded string.
//maps are encoded in the rails style (key[key2][key2]=value)
// TODO: we should sort the keynames so ordering is consistent and then this
// can be used a cache key
func (this *DynMap) MarshalURL() (string, error) {
	vals := &url.Values{}
	for key, value := range this.Map {
		err := this.urlEncode(vals, key, value)
		if err != nil {
			return "", err
		}
	}
	return vals.Encode(), nil
}

// Unmarshals a url encoded string.
// will also parse rails style maps in the form key[key1][key2]=val\
func (this *DynMap) UnmarshalURL(urlstring string) error {
	//TODO: split on ?
	values, err := url.ParseQuery(urlstring)
	if err != nil {
		return err
	}
	for k := range values {
		var v = values[k]
		key := strings.Replace(k, "[", ".", -1)
		key = strings.Replace(key, "]", "", -1)

		if len(v) == 1 {
			this.PutWithDot(key, v[0])
		} else {
			this.PutWithDot(key, v)
		}
	}
	return nil
}

//adds the requested value to the Values
func (this *DynMap) urlEncode(vals *url.Values, key string, value interface{}) error {

	if DynMapConvertable(value) {
		mp, ok := ToDynMap(value)
		if !ok {
			return fmt.Errorf("Unable to convert %s", mp)
		}
		for k, v := range mp.Map {
			//encode in rails style key[key2]=value
			this.urlEncode(vals, fmt.Sprintf("%s[%s]", key, k), v)
		}
		return nil
	}
	switch v := value.(type) {
	case []interface{}:
		for _, tmp := range v {
			this.urlEncode(vals, key, tmp)
		}
		return nil
	}
	vals.Add(key, ToString(value))
	return nil
}

func (this *DynMap) MarshalJSON() ([]byte, error) {
	bytes, err := json.Marshal(this.Map)
	return bytes, err
}

func (this *DynMap) UnmarshalJSON(bytes []byte) error {
	return json.Unmarshal(bytes, &this.Map)
}

// Gets the value at the specified key as an int64.  returns -1,false if value not available or is not convertable
func (this *DynMap) GetInt64(key string) (int64, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return -1, ok
	}
	val, err := ToInt64(tmp)
	if err == nil {
		return val, true
	}
	return -1, false
}

func (this *DynMap) MustInt64(key string, def int64) int64 {
	v, ok := this.GetInt64(key)
	if ok {
		return v
	}
	return def
}

func (this *DynMap) MustInt(key string, def int) int {
	v, ok := this.GetInt(key)
	if ok {
		return v
	}
	return def
}

func (this *DynMap) GetInt(key string) (int, bool) {
	v, ok := this.GetInt64(key)
	if !ok {
		return -1, ok
	}
	return int(v), true
}

//
// Gets a string representation of the value at key
//
func (this *DynMap) GetString(key string) (string, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return ToString(tmp), ok
	}
	return ToString(tmp), true
}

// gets a string. if string is not available in the map, then the default
//is returned
func (this *DynMap) MustString(key string, def string) string {
	tmp, ok := this.GetString(key)
	if !ok {
		return def
	}
	return tmp
}

func (this *DynMap) GetTime(key string) (time.Time, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return time.Now(), false
	}
	t, err := ToTime(tmp)
	if err != nil {
		return time.Now(), false
	}
	return t, true
}

func (this *DynMap) MustTime(key string, def time.Time) time.Time {
	tmp, ok := this.GetTime(key)
	if !ok {
		return def
	}
	return tmp
}

func (this *DynMap) GetBool(key string) (bool, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return false, ok
	}
	b, err := ToBool(tmp)
	if err != nil {
		return false, false
	}
	return b, true
}

func (this *DynMap) MustBool(key string, def bool) bool {
	tmp, ok := this.GetBool(key)
	if !ok {
		return def
	}
	return tmp
}

//Gets a dynmap from the requested.
// This will update the value in the map if the
// value was not already a dynmap.
func (this *DynMap) GetDynMap(key string) (*DynMap, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return nil, ok
	}
	mp, ok := ToDynMap(tmp)
	return mp, ok
}

func (this *DynMap) MustDynMap(key string, def *DynMap) *DynMap {
	tmp, ok := this.GetDynMap(key)
	if !ok {
		return def
	}
	return tmp
}

// gets a slice of dynmaps
func (this *DynMap) GetDynMapSlice(key string) ([]*DynMap, bool) {
	lst, ok := this.Get(key)
	if !ok {
		return nil, false
	}
	switch v := lst.(type) {
	case []*DynMap:
		return v, true
	case []interface{}:
		retlist := make([]*DynMap, 0)
		for _, tmp := range v {
			in, ok := ToDynMap(tmp)
			if !ok {
				return nil, false
			}
			retlist = append(retlist, in)
		}
		return retlist, true
	}
	return nil, false
}

//Returns a slice of ints
func (this *DynMap) GetIntSlice(key string) ([]int, bool) {
	lst, ok := this.Get(key)
	if !ok {
		return nil, false
	}
	switch v := lst.(type) {
	case []int:
		return v, true
	case []interface{}:
		retlist := make([]int, 0)
		for _, tmp := range v {
			in, err := ToInt(tmp)
			if err != nil {
				return nil, false
			}
			retlist = append(retlist, in)
		}
		return retlist, true
	}
	return nil, false
}

//gets a slice of ints.  if the value is a string it will
//split by the requested delimiter
func (this *DynMap) GetIntSliceSplit(key, delim string) ([]int, bool) {
	lst, ok := this.Get(key)
	if !ok {
		return nil, false
	}
	switch v := lst.(type) {
	case string:
		retlist := make([]int, 0)
		for _, tmp := range strings.Split(v, delim) {
			in, err := ToInt(tmp)
			if err != nil {
				return nil, false
			}
			retlist = append(retlist, in)
		}
		return retlist, true
	}
	ret, ok := this.GetIntSlice(key)
	return ret, ok
}

// Adds the item to a slice
func (this *DynMap) AddToSlice(key string, mp interface{}) error {
	this.PutIfAbsent(key, make([]interface{}, 0))
	lst, _ := this.Get(key)
	switch v := lst.(type) {
	case []interface{}:
		v = append(v, mp)
		this.Put(key, v)
	}
	return nil
}

// puts all the values from the passed in map into this dynmap
// the passed in map must be convertable to a DynMap via ToDynMap.
// returns false if the passed value is not convertable to dynmap
func (this *DynMap) PutAll(mp interface{}) bool {
	dynmap, ok := ToDynMap(mp)
	if !ok {
		return false
	}
	for k, v := range dynmap.Map {
		this.Put(k, v)
	}
	return true
}

//
// Puts the value into the map if and only if no value exists at the
// specified key.
// This does not honor the dot operator on insert.
func (this *DynMap) PutIfAbsent(key string, value interface{}) (interface{}, bool) {
	v, ok := this.Get(key)
	if ok {
		return v, false
	}
	this.Put(key, value)
	return value, true
}

//
// Same as PutIfAbsent but honors the dot operator
//
func (this *DynMap) PutIfAbsentWithDot(key string, value interface{}) (interface{}, bool) {
	v, ok := this.Get(key)
	if ok {
		return v, false
	}
	this.PutWithDot(key, value)
	return value, true
}

//
// Put's a value into the map
//
func (this *DynMap) Put(key string, value interface{}) {
	this.Map[key] = value
}

//
// puts the value into the map, honoring the dot operator.
// so PutWithDot("map1.map2.value", 100)
// would result in:
// {
//   map1 : { map2 : { value: 100 }}
//
// }
func (this *DynMap) PutWithDot(key string, value interface{}) error {
	splitStr := strings.Split(key, ".")
	if len(splitStr) == 1 {
		this.Put(key, value)
		return nil
	}
	mapKeys := splitStr[:(len(splitStr) - 1)]
	var mp = this.Map
	for _, k := range mapKeys {
		tmp, o := mp[k]
		if !o {
			//create a new map and insert
			newmap := make(map[string]interface{})
			mp[k] = newmap
			mp = newmap
		} else {
			mp, o = ToMap(tmp)
			if !o {
				//error
				return errors.New("Error, value at key was not a map")
			}
		}
	}
	mp[splitStr[len(splitStr)-1]] = value
	return nil
}

func (this *DynMap) Exists(key string) bool {
	_, ok := this.Get(key)
	return ok
}

//Remove a mapping
//TODO: this will need to honor the dot operator!
func (this *DynMap) Remove(key string) (interface{}, bool) {
	val, ok := this.Map[key]
	if ok {
		delete(this.Map, key)
		return val, true
	}
	//TODO: dot op..
	return val, false

}

//
// Get's the value.  will honor the dot operator if needed.
// key = 'map.map2'
// will first attempt to matche the literal key 'map.map2'
// if no value is present it will look for a sub map at key 'map'
//
func (this *DynMap) Get(key string) (interface{}, bool) {
	val, ok := this.Map[key]
	if ok {
		return val, true
	}
	//look for dot operator.
	splitStr := strings.Split(key, ".")
	if len(splitStr) == 1 {
		return val, false
	}

	var mp = this.Map
	for index, k := range splitStr {
		tmp, o := mp[k]
		if !o {
			return val, ok
		}

		if index == (len(splitStr) - 1) {
			return tmp, o
		} else {
			mp, o = ToMap(tmp)
			if !o {
				return val, ok
			}
		}
	}
	return val, ok
}
<commit_msg>add New dynmap method<commit_after>package dynmap

import (
	"strings"
	// "log"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"time"
)

//Dont make this a map type, since we want the option of
//extending this and adding members.
type DynMap struct {
	Map map[string]interface{}
}

type DynMaper interface {
	ToDynMap() *DynMap
}

// Creates a new dynmap
func New() *DynMap {
	return &DynMap{make(map[string]interface{})}
}

// Deprecated.  use New() instead
func NewDynMap() *DynMap {
	return &DynMap{make(map[string]interface{})}
}

// Recursively converts this to a regular go map.
// (will convert any sub maps)
func (this *DynMap) ToMap() map[string]interface{} {
	mp := make(map[string]interface{})
	for k, v := range(this.Map) {
		submp, ok := ToDynMap(this.Map[k])
		if ok {
			v = submp.ToMap()
		}
		mp[k] = v
	}
	return mp
}

// Returns self. Here so that we satisfy the DynMaper interface
func (this *DynMap) ToDynMap() *DynMap {
	return this
}

//encodes this map into a url encoded string.
//maps are encoded in the rails style (key[key2][key2]=value)
// TODO: we should sort the keynames so ordering is consistent and then this
// can be used a cache key
func (this *DynMap) MarshalURL() (string, error) {
	vals := &url.Values{}
	for key, value := range this.Map {
		err := this.urlEncode(vals, key, value)
		if err != nil {
			return "", err
		}
	}
	return vals.Encode(), nil
}

// Unmarshals a url encoded string.
// will also parse rails style maps in the form key[key1][key2]=val\
func (this *DynMap) UnmarshalURL(urlstring string) error {
	//TODO: split on ?
	values, err := url.ParseQuery(urlstring)
	if err != nil {
		return err
	}
	for k := range values {
		var v = values[k]
		key := strings.Replace(k, "[", ".", -1)
		key = strings.Replace(key, "]", "", -1)

		if len(v) == 1 {
			this.PutWithDot(key, v[0])
		} else {
			this.PutWithDot(key, v)
		}
	}
	return nil
}

//adds the requested value to the Values
func (this *DynMap) urlEncode(vals *url.Values, key string, value interface{}) error {

	if DynMapConvertable(value) {
		mp, ok := ToDynMap(value)
		if !ok {
			return fmt.Errorf("Unable to convert %s", mp)
		}
		for k, v := range mp.Map {
			//encode in rails style key[key2]=value
			this.urlEncode(vals, fmt.Sprintf("%s[%s]", key, k), v)
		}
		return nil
	}
	switch v := value.(type) {
	case []interface{}:
		for _, tmp := range v {
			this.urlEncode(vals, key, tmp)
		}
		return nil
	}
	vals.Add(key, ToString(value))
	return nil
}

func (this *DynMap) MarshalJSON() ([]byte, error) {
	bytes, err := json.Marshal(this.Map)
	return bytes, err
}

func (this *DynMap) UnmarshalJSON(bytes []byte) error {
	return json.Unmarshal(bytes, &this.Map)
}

// Gets the value at the specified key as an int64.  returns -1,false if value not available or is not convertable
func (this *DynMap) GetInt64(key string) (int64, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return -1, ok
	}
	val, err := ToInt64(tmp)
	if err == nil {
		return val, true
	}
	return -1, false
}

func (this *DynMap) MustInt64(key string, def int64) int64 {
	v, ok := this.GetInt64(key)
	if ok {
		return v
	}
	return def
}

func (this *DynMap) MustInt(key string, def int) int {
	v, ok := this.GetInt(key)
	if ok {
		return v
	}
	return def
}

func (this *DynMap) GetInt(key string) (int, bool) {
	v, ok := this.GetInt64(key)
	if !ok {
		return -1, ok
	}
	return int(v), true
}

//
// Gets a string representation of the value at key
//
func (this *DynMap) GetString(key string) (string, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return ToString(tmp), ok
	}
	return ToString(tmp), true
}

// gets a string. if string is not available in the map, then the default
//is returned
func (this *DynMap) MustString(key string, def string) string {
	tmp, ok := this.GetString(key)
	if !ok {
		return def
	}
	return tmp
}

func (this *DynMap) GetTime(key string) (time.Time, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return time.Now(), false
	}
	t, err := ToTime(tmp)
	if err != nil {
		return time.Now(), false
	}
	return t, true
}

func (this *DynMap) MustTime(key string, def time.Time) time.Time {
	tmp, ok := this.GetTime(key)
	if !ok {
		return def
	}
	return tmp
}

func (this *DynMap) GetBool(key string) (bool, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return false, ok
	}
	b, err := ToBool(tmp)
	if err != nil {
		return false, false
	}
	return b, true
}

func (this *DynMap) MustBool(key string, def bool) bool {
	tmp, ok := this.GetBool(key)
	if !ok {
		return def
	}
	return tmp
}

//Gets a dynmap from the requested.
// This will update the value in the map if the
// value was not already a dynmap.
func (this *DynMap) GetDynMap(key string) (*DynMap, bool) {
	tmp, ok := this.Get(key)
	if !ok {
		return nil, ok
	}
	mp, ok := ToDynMap(tmp)
	return mp, ok
}

func (this *DynMap) MustDynMap(key string, def *DynMap) *DynMap {
	tmp, ok := this.GetDynMap(key)
	if !ok {
		return def
	}
	return tmp
}

// gets a slice of dynmaps
func (this *DynMap) GetDynMapSlice(key string) ([]*DynMap, bool) {
	lst, ok := this.Get(key)
	if !ok {
		return nil, false
	}
	switch v := lst.(type) {
	case []*DynMap:
		return v, true
	case []interface{}:
		retlist := make([]*DynMap, 0)
		for _, tmp := range v {
			in, ok := ToDynMap(tmp)
			if !ok {
				return nil, false
			}
			retlist = append(retlist, in)
		}
		return retlist, true
	}
	return nil, false
}

//Returns a slice of ints
func (this *DynMap) GetIntSlice(key string) ([]int, bool) {
	lst, ok := this.Get(key)
	if !ok {
		return nil, false
	}
	switch v := lst.(type) {
	case []int:
		return v, true
	case []interface{}:
		retlist := make([]int, 0)
		for _, tmp := range v {
			in, err := ToInt(tmp)
			if err != nil {
				return nil, false
			}
			retlist = append(retlist, in)
		}
		return retlist, true
	}
	return nil, false
}

//gets a slice of ints.  if the value is a string it will
//split by the requested delimiter
func (this *DynMap) GetIntSliceSplit(key, delim string) ([]int, bool) {
	lst, ok := this.Get(key)
	if !ok {
		return nil, false
	}
	switch v := lst.(type) {
	case string:
		retlist := make([]int, 0)
		for _, tmp := range strings.Split(v, delim) {
			in, err := ToInt(tmp)
			if err != nil {
				return nil, false
			}
			retlist = append(retlist, in)
		}
		return retlist, true
	}
	ret, ok := this.GetIntSlice(key)
	return ret, ok
}

// Adds the item to a slice
func (this *DynMap) AddToSlice(key string, mp interface{}) error {
	this.PutIfAbsent(key, make([]interface{}, 0))
	lst, _ := this.Get(key)
	switch v := lst.(type) {
	case []interface{}:
		v = append(v, mp)
		this.Put(key, v)
	}
	return nil
}

// puts all the values from the passed in map into this dynmap
// the passed in map must be convertable to a DynMap via ToDynMap.
// returns false if the passed value is not convertable to dynmap
func (this *DynMap) PutAll(mp interface{}) bool {
	dynmap, ok := ToDynMap(mp)
	if !ok {
		return false
	}
	for k, v := range dynmap.Map {
		this.Put(k, v)
	}
	return true
}

//
// Puts the value into the map if and only if no value exists at the
// specified key.
// This does not honor the dot operator on insert.
func (this *DynMap) PutIfAbsent(key string, value interface{}) (interface{}, bool) {
	v, ok := this.Get(key)
	if ok {
		return v, false
	}
	this.Put(key, value)
	return value, true
}

//
// Same as PutIfAbsent but honors the dot operator
//
func (this *DynMap) PutIfAbsentWithDot(key string, value interface{}) (interface{}, bool) {
	v, ok := this.Get(key)
	if ok {
		return v, false
	}
	this.PutWithDot(key, value)
	return value, true
}

//
// Put's a value into the map
//
func (this *DynMap) Put(key string, value interface{}) {
	this.Map[key] = value
}

//
// puts the value into the map, honoring the dot operator.
// so PutWithDot("map1.map2.value", 100)
// would result in:
// {
//   map1 : { map2 : { value: 100 }}
//
// }
func (this *DynMap) PutWithDot(key string, value interface{}) error {
	splitStr := strings.Split(key, ".")
	if len(splitStr) == 1 {
		this.Put(key, value)
		return nil
	}
	mapKeys := splitStr[:(len(splitStr) - 1)]
	var mp = this.Map
	for _, k := range mapKeys {
		tmp, o := mp[k]
		if !o {
			//create a new map and insert
			newmap := make(map[string]interface{})
			mp[k] = newmap
			mp = newmap
		} else {
			mp, o = ToMap(tmp)
			if !o {
				//error
				return errors.New("Error, value at key was not a map")
			}
		}
	}
	mp[splitStr[len(splitStr)-1]] = value
	return nil
}

func (this *DynMap) Exists(key string) bool {
	_, ok := this.Get(key)
	return ok
}

//Remove a mapping
//TODO: this will need to honor the dot operator!
func (this *DynMap) Remove(key string) (interface{}, bool) {
	val, ok := this.Map[key]
	if ok {
		delete(this.Map, key)
		return val, true
	}
	//TODO: dot op..
	return val, false

}

//
// Get's the value.  will honor the dot operator if needed.
// key = 'map.map2'
// will first attempt to matche the literal key 'map.map2'
// if no value is present it will look for a sub map at key 'map'
//
func (this *DynMap) Get(key string) (interface{}, bool) {
	val, ok := this.Map[key]
	if ok {
		return val, true
	}
	//look for dot operator.
	splitStr := strings.Split(key, ".")
	if len(splitStr) == 1 {
		return val, false
	}

	var mp = this.Map
	for index, k := range splitStr {
		tmp, o := mp[k]
		if !o {
			return val, ok
		}

		if index == (len(splitStr) - 1) {
			return tmp, o
		} else {
			mp, o = ToMap(tmp)
			if !o {
				return val, ok
			}
		}
	}
	return val, ok
}
<|endoftext|>


<commit_before>// Package mucp transparently forwards the incoming request using a go-micro client.
package mucp

import (
	"context"
	"fmt"
	"io"
	"strings"
	"sync"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/codec/bytes"
	"github.com/micro/go-micro/config/options"
	"github.com/micro/go-micro/proxy"
	"github.com/micro/go-micro/router"
	"github.com/micro/go-micro/server"
)

// Proxy will transparently proxy requests to an endpoint.
// If no endpoint is specified it will call a service using the client.
type Proxy struct {
	// embed options
	options.Options

	// Endpoint specifies the fixed service endpoint to call.
	Endpoint string

	// The client to use for outbound requests
	Client client.Client

	// The router for routes
	Router router.Router

	// A fib of routes service:address
	sync.RWMutex
	Routes map[string]map[uint64]router.Route

	// The channel to monitor watcher errors
	errChan chan error
}

// read client request and write to server
func readLoop(r server.Request, s client.Stream) error {
	// request to backend server
	req := s.Request()

	for {
		// get data from client
		//  no need to decode it
		body, err := r.Read()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		// get the header from client
		hdr := r.Header()
		msg := &codec.Message{
			Type:   codec.Request,
			Header: hdr,
			Body:   body,
		}

		// write the raw request
		err = req.Codec().Write(msg, nil)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

// toNodes returns a list of node addresses from given routes
func toNodes(routes map[uint64]router.Route) []string {
	var nodes []string
	for _, node := range routes {
		address := node.Address
		if len(node.Gateway) > 0 {
			address = node.Gateway
		}
		nodes = append(nodes, address)
	}
	return nodes
}

func (p *Proxy) getRoute(service string) ([]string, error) {
	// lookup the route cache first
	p.Lock()
	routes, ok := p.Routes[service]
	if ok {
		p.Unlock()
		return toNodes(routes), nil
	}
	p.Routes[service] = make(map[uint64]router.Route)
	p.Unlock()

	// lookup the routes in the router
	results, err := p.Router.Lookup(router.NewQuery(router.QueryService(service)))
	if err != nil {
		// check the status of the router
		if status := p.Router.Status(); status.Code == router.Error {
			return nil, status.Error
		}
		// otherwise return the error
		return nil, err
	}

	// update the proxy cache
	p.Lock()
	for _, route := range results {
		p.Routes[service][route.Hash()] = route
	}
	routes = p.Routes[service]
	p.Unlock()

	return toNodes(routes), nil
}

// manageRouteCache applies action on a given route to Proxy route cache
func (p *Proxy) manageRouteCache(route router.Route, action string) error {
	switch action {
	case "create", "update":
		if _, ok := p.Routes[route.Service]; !ok {
			p.Routes[route.Service] = make(map[uint64]router.Route)
		}
		p.Routes[route.Service][route.Hash()] = route
	case "delete":
		if _, ok := p.Routes[route.Service]; !ok {
			return fmt.Errorf("route not found")
		}
		delete(p.Routes[route.Service], route.Hash())
	default:
		return fmt.Errorf("unknown action: %s", action)
	}

	return nil
}

// watchRoutes watches service routes and updates proxy cache
func (p *Proxy) watchRoutes() {
	// this is safe to do as the only way watchRoutes returns is
	// when some error is written into error channel - we want to bail then
	defer close(p.errChan)

	// route watcher
	w, err := p.Router.Watch()
	if err != nil {
		p.errChan <- err
		return
	}

	for {
		event, err := w.Next()
		if err != nil {
			p.errChan <- err
			return
		}

		p.Lock()
		if err := p.manageRouteCache(event.Route, fmt.Sprintf("%s", event.Type)); err != nil {
			// TODO: should we bail here?
			p.Unlock()
			continue
		}
		p.Unlock()
	}
}

// ServeRequest honours the server.Router interface
func (p *Proxy) ServeRequest(ctx context.Context, req server.Request, rsp server.Response) error {
	// service name
	service := req.Service()
	endpoint := req.Endpoint()
	var addresses []string

	// call a specific backend endpoint either by name or address
	if len(p.Endpoint) > 0 {
		// address:port
		if parts := strings.Split(p.Endpoint, ":"); len(parts) > 1 {
			addresses = []string{p.Endpoint}
		} else {
			// get route for endpoint from router
			addr, err := p.getRoute(p.Endpoint)
			if err != nil {
				return err
			}
			// set the address
			addresses = addr
			// set the name
			service = p.Endpoint
		}
	} else {
		// no endpoint was specified just lookup the route
		// get route for endpoint from router
		addr, err := p.getRoute(service)
		if err != nil {
			return err
		}
		addresses = addr
	}

	var opts []client.CallOption

	// set address if available
	if len(addresses) > 0 {
		opts = append(opts, client.WithAddress(addresses...))
	}

	// read initial request
	body, err := req.Read()
	if err != nil {
		return err
	}

	// create new request with raw bytes body
	creq := p.Client.NewRequest(service, endpoint, &bytes.Frame{body}, client.WithContentType(req.ContentType()))

	// create new stream
	stream, err := p.Client.Stream(ctx, creq, opts...)
	if err != nil {
		return err
	}
	defer stream.Close()

	// create client request read loop
	go readLoop(req, stream)

	// get raw response
	resp := stream.Response()

	// route watcher error
	var watchErr error

	// create server response write loop
	for {
		select {
		case err := <-p.errChan:
			if err != nil {
				watchErr = err
			}
			return watchErr
		default:
			// read backend response body
			body, err := resp.Read()
			if err == io.EOF {
				return nil
			} else if err != nil {
				return err
			}

			// read backend response header
			hdr := resp.Header()

			// write raw response header to client
			rsp.WriteHeader(hdr)

			// write raw response body to client
			err = rsp.Write(body)
			if err == io.EOF {
				return nil
			} else if err != nil {
				return err
			}
		}
	}

	return nil
}

// NewSingleHostProxy returns a proxy which sends requests to a single backend
func NewSingleHostProxy(endpoint string) *Proxy {
	return &Proxy{
		Options:  options.NewOptions(),
		Endpoint: endpoint,
	}
}

// NewProxy returns a new proxy which will route based on mucp headers
func NewProxy(opts ...options.Option) proxy.Proxy {
	p := new(Proxy)
	p.Options = options.NewOptions(opts...)
	p.Options.Init(options.WithString("mucp"))

	// get endpoint
	ep, ok := p.Options.Values().Get("proxy.endpoint")
	if ok {
		p.Endpoint = ep.(string)
	}

	// get client
	c, ok := p.Options.Values().Get("proxy.client")
	if ok {
		p.Client = c.(client.Client)
	}

	// set the default client
	if p.Client == nil {
		p.Client = client.DefaultClient
	}

	// get router
	r, ok := p.Options.Values().Get("proxy.router")
	if ok {
		p.Router = r.(router.Router)
	}

	// create default router and start it
	if p.Router == nil {
		p.Router = router.DefaultRouter
	}

	// routes cache
	p.Routes = make(map[string]map[uint64]router.Route)

	// watch router service routes
	p.errChan = make(chan error, 1)
	go p.watchRoutes()

	return p
}
<commit_msg>only continue to stream when its a stream<commit_after>// Package mucp transparently forwards the incoming request using a go-micro client.
package mucp

import (
	"context"
	"fmt"
	"io"
	"strings"
	"sync"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/codec"
	"github.com/micro/go-micro/codec/bytes"
	"github.com/micro/go-micro/config/options"
	"github.com/micro/go-micro/proxy"
	"github.com/micro/go-micro/router"
	"github.com/micro/go-micro/server"
)

// Proxy will transparently proxy requests to an endpoint.
// If no endpoint is specified it will call a service using the client.
type Proxy struct {
	// embed options
	options.Options

	// Endpoint specifies the fixed service endpoint to call.
	Endpoint string

	// The client to use for outbound requests
	Client client.Client

	// The router for routes
	Router router.Router

	// A fib of routes service:address
	sync.RWMutex
	Routes map[string]map[uint64]router.Route

	// The channel to monitor watcher errors
	errChan chan error
}

// read client request and write to server
func readLoop(r server.Request, s client.Stream) error {
	// we don't loop unless its a stream
	if !r.Stream() {
		return nil
	}

	// request to backend server
	req := s.Request()

	for {
		// get data from client
		//  no need to decode it
		body, err := r.Read()
		if err == io.EOF {
			return nil
		}

		if err != nil {
			return err
		}

		// get the header from client
		hdr := r.Header()
		msg := &codec.Message{
			Type:   codec.Request,
			Header: hdr,
			Body:   body,
		}

		// write the raw request
		err = req.Codec().Write(msg, nil)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

// toNodes returns a list of node addresses from given routes
func toNodes(routes map[uint64]router.Route) []string {
	var nodes []string
	for _, node := range routes {
		address := node.Address
		if len(node.Gateway) > 0 {
			address = node.Gateway
		}
		nodes = append(nodes, address)
	}
	return nodes
}

func (p *Proxy) getRoute(service string) ([]string, error) {
	// lookup the route cache first
	p.Lock()
	routes, ok := p.Routes[service]
	if ok {
		p.Unlock()
		return toNodes(routes), nil
	}
	p.Routes[service] = make(map[uint64]router.Route)
	p.Unlock()

	// lookup the routes in the router
	results, err := p.Router.Lookup(router.NewQuery(router.QueryService(service)))
	if err != nil {
		// check the status of the router
		if status := p.Router.Status(); status.Code == router.Error {
			return nil, status.Error
		}
		// otherwise return the error
		return nil, err
	}

	// update the proxy cache
	p.Lock()
	for _, route := range results {
		p.Routes[service][route.Hash()] = route
	}
	routes = p.Routes[service]
	p.Unlock()

	return toNodes(routes), nil
}

// manageRouteCache applies action on a given route to Proxy route cache
func (p *Proxy) manageRouteCache(route router.Route, action string) error {
	switch action {
	case "create", "update":
		if _, ok := p.Routes[route.Service]; !ok {
			p.Routes[route.Service] = make(map[uint64]router.Route)
		}
		p.Routes[route.Service][route.Hash()] = route
	case "delete":
		if _, ok := p.Routes[route.Service]; !ok {
			return fmt.Errorf("route not found")
		}
		delete(p.Routes[route.Service], route.Hash())
	default:
		return fmt.Errorf("unknown action: %s", action)
	}

	return nil
}

// watchRoutes watches service routes and updates proxy cache
func (p *Proxy) watchRoutes() {
	// this is safe to do as the only way watchRoutes returns is
	// when some error is written into error channel - we want to bail then
	defer close(p.errChan)

	// route watcher
	w, err := p.Router.Watch()
	if err != nil {
		p.errChan <- err
		return
	}

	for {
		event, err := w.Next()
		if err != nil {
			p.errChan <- err
			return
		}

		p.Lock()
		if err := p.manageRouteCache(event.Route, fmt.Sprintf("%s", event.Type)); err != nil {
			// TODO: should we bail here?
			p.Unlock()
			continue
		}
		p.Unlock()
	}
}

// ServeRequest honours the server.Router interface
func (p *Proxy) ServeRequest(ctx context.Context, req server.Request, rsp server.Response) error {
	// service name
	service := req.Service()
	endpoint := req.Endpoint()
	var addresses []string

	// call a specific backend endpoint either by name or address
	if len(p.Endpoint) > 0 {
		// address:port
		if parts := strings.Split(p.Endpoint, ":"); len(parts) > 1 {
			addresses = []string{p.Endpoint}
		} else {
			// get route for endpoint from router
			addr, err := p.getRoute(p.Endpoint)
			if err != nil {
				return err
			}
			// set the address
			addresses = addr
			// set the name
			service = p.Endpoint
		}
	} else {
		// no endpoint was specified just lookup the route
		// get route for endpoint from router
		addr, err := p.getRoute(service)
		if err != nil {
			return err
		}
		addresses = addr
	}

	var opts []client.CallOption

	// set address if available
	if len(addresses) > 0 {
		opts = append(opts, client.WithAddress(addresses...))
	}

	// read initial request
	body, err := req.Read()
	if err != nil {
		return err
	}

	// create new request with raw bytes body
	creq := p.Client.NewRequest(service, endpoint, &bytes.Frame{body}, client.WithContentType(req.ContentType()))

	// create new stream
	stream, err := p.Client.Stream(ctx, creq, opts...)
	if err != nil {
		return err
	}
	defer stream.Close()

	// create client request read loop
	go readLoop(req, stream)

	// get raw response
	resp := stream.Response()

	// route watcher error
	var watchErr error

	// create server response write loop
	for {
		select {
		case err := <-p.errChan:
			if err != nil {
				watchErr = err
			}
			return watchErr
		default:
			// read backend response body
			body, err := resp.Read()
			if err == io.EOF {
				return nil
			} else if err != nil {
				return err
			}

			// read backend response header
			hdr := resp.Header()

			// write raw response header to client
			rsp.WriteHeader(hdr)

			// write raw response body to client
			err = rsp.Write(body)
			if err == io.EOF {
				return nil
			} else if err != nil {
				return err
			}

			// we don't continue unless its a stream
			if !req.Stream() {
				return nil
			}
		}
	}

	return nil
}

// NewSingleHostProxy returns a proxy which sends requests to a single backend
func NewSingleHostProxy(endpoint string) *Proxy {
	return &Proxy{
		Options:  options.NewOptions(),
		Endpoint: endpoint,
	}
}

// NewProxy returns a new proxy which will route based on mucp headers
func NewProxy(opts ...options.Option) proxy.Proxy {
	p := new(Proxy)
	p.Options = options.NewOptions(opts...)
	p.Options.Init(options.WithString("mucp"))

	// get endpoint
	ep, ok := p.Options.Values().Get("proxy.endpoint")
	if ok {
		p.Endpoint = ep.(string)
	}

	// get client
	c, ok := p.Options.Values().Get("proxy.client")
	if ok {
		p.Client = c.(client.Client)
	}

	// set the default client
	if p.Client == nil {
		p.Client = client.DefaultClient
	}

	// get router
	r, ok := p.Options.Values().Get("proxy.router")
	if ok {
		p.Router = r.(router.Router)
	}

	// create default router and start it
	if p.Router == nil {
		p.Router = router.DefaultRouter
	}

	// routes cache
	p.Routes = make(map[string]map[uint64]router.Route)

	// watch router service routes
	p.errChan = make(chan error, 1)
	go p.watchRoutes()

	return p
}
<|endoftext|>


<commit_before>package socks

import (
	"net"

	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

const (
	bufferSize = 2 * 1024
)

var udpAddress v2net.Address

func (server *SocksServer) ListenUDP(port uint16) error {
	addr := &net.UDPAddr{
		IP:   net.IP{0, 0, 0, 0},
		Port: int(port),
		Zone: "",
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
		return err
	}
	udpAddress = v2net.IPAddress([]byte{127, 0, 0, 1}, port)

	go server.AcceptPackets(conn)
	return nil
}

func (server *SocksServer) getUDPAddr() v2net.Address {
	return udpAddress
}

func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
	for {
		buffer := make([]byte, bufferSize)
		nBytes, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Error("Socks failed to read UDP packets: %v", err)
			return err
		}
		request, err := protocol.ReadUDPRequest(buffer[:nBytes])
		if err != nil {
			log.Error("Socks failed to parse UDP request: %v", err)
			return err
		}
		if request.Fragment != 0 {
			// TODO handle fragments
			continue
		}

		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
		go server.handlePacket(conn, udpPacket, addr, request)
	}
}

func (server *SocksServer) handlePacket(conn *net.UDPConn, packet v2net.Packet, clientAddr *net.UDPAddr, request protocol.Socks5UDPRequest) {
	ray := server.vPoint.DispatchToOutbound(packet)
	close(ray.InboundInput())

	if data, ok := <-ray.InboundOutput(); ok {
		request.Data = data
		udpMessage := request.Bytes(nil)
		nBytes, err := conn.WriteToUDP(udpMessage, clientAddr)
		if err != nil {
			log.Error("Socks failed to write UDP message (%d bytes) to %s: %v", nBytes, clientAddr.String(), err)
		}
	}
}
<commit_msg>Fix UDP addr<commit_after>package socks

import (
	"net"

	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

const (
	bufferSize = 2 * 1024
)

var udpAddress v2net.Address

func (server *SocksServer) ListenUDP(port uint16) error {
	addr := &net.UDPAddr{
		IP:   net.IP{0, 0, 0, 0},
		Port: int(port),
		Zone: "",
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
		return err
	}
	udpAddress = v2net.IPAddress(conn.LocalAddr().(*net.UDPAddr).IP, uint16(conn.LocalAddr().(*net.UDPAddr).Port))

	go server.AcceptPackets(conn)
	return nil
}

func (server *SocksServer) getUDPAddr() v2net.Address {
	return udpAddress
}

func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
	for {
		buffer := make([]byte, bufferSize)
		nBytes, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Error("Socks failed to read UDP packets: %v", err)
			return err
		}
		request, err := protocol.ReadUDPRequest(buffer[:nBytes])
		if err != nil {
			log.Error("Socks failed to parse UDP request: %v", err)
			return err
		}
		if request.Fragment != 0 {
			// TODO handle fragments
			continue
		}

		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
		go server.handlePacket(conn, udpPacket, addr, request)
	}
}

func (server *SocksServer) handlePacket(conn *net.UDPConn, packet v2net.Packet, clientAddr *net.UDPAddr, request protocol.Socks5UDPRequest) {
	ray := server.vPoint.DispatchToOutbound(packet)
	close(ray.InboundInput())

	if data, ok := <-ray.InboundOutput(); ok {
		request.Data = data
		udpMessage := request.Bytes(nil)
		nBytes, err := conn.WriteToUDP(udpMessage, clientAddr)
		if err != nil {
			log.Error("Socks failed to write UDP message (%d bytes) to %s: %v", nBytes, clientAddr.String(), err)
		}
	}
}
<|endoftext|>