<commit_before>package main

import (
	"fmt"
	"github.com/dustin/go-humanize"
	"github.com/gorilla/mux"
	"github.com/libgit2/git2go"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"path"
	"path/filepath"
	"strings"
	"time"
)

type TreeEntry struct {
	// extend with a pointer to the parent tree
	DirPath string

	*git.TreeEntry
}
type Pathable interface {
	Path() string
}

type Repo struct {
	Name        string
	Path        string
	Description string
	_Repo       *git.Repository
}

type Commit struct {
	_Commit *git.Commit
	Message string
	Author  string
	Hash    string
	Date    time.Time
}

type Diff struct {
	_CommitA    *git.Commit
	_CommitB    *git.Commit
	CommitAHash string
	CommitBHash string
	Stats       string
	Patches     []string
}

func (c Commit) Path() string {
	return path.Join("commit", c.Hash)
}

func (t TreeEntry) Path() string {
	if t.Type == git.ObjectBlob {
		return filepath.Join("/", "blob", t.DirPath, t.Name)
	} else if t.Type == git.ObjectTree {
		if t.DirPath == "" {
			return "/"
		} else {
			return filepath.Join("/", "tree", t.DirPath, t.Name)
		}
	}
	log.Fatal("unknown type tree entry type ", t.Type)
	return ""
}

// TODO: get commit log for an arbitrary branch
func getCommitLog(repo *git.Repository, obj *git.Object) []Commit {
	r, err := repo.Walk()
	if err != nil {
		log.Print("failed to walk repo: ", err)
	}

	r.Push(obj.Id())
	r.Sorting(git.SortTime)
	r.SimplifyFirstParent()

	id := &(git.Oid{})

	var commits []Commit
	for r.Next(id) == nil {
		g, _ := repo.LookupCommit(id)
		c := Commit{g, g.Message(), g.Committer().Name, g.Id().String(), g.Committer().When}
		commits = append(commits, c)
	}
	return commits
}

func getCommitDiff(repo *git.Repository, commit *git.Commit) Diff {
	// TODO check for multiple parent commits
	p, _ := commit.Parent(0).Tree()
	c, _ := commit.Tree()
	o, _ := git.DefaultDiffOptions()
	diff, _ := repo.DiffTreeToTree(p, c, &o)
	defer diff.Free()

	stats, _ := diff.Stats()
	statsStr, _ := stats.String(git.DiffStatsFull, 80)

	r := Diff{
		_CommitA:    commit.Parent(0),
		_CommitB:    commit,
		CommitAHash: p.Id().String(),
		CommitBHash: c.Id().String(),
		Stats:       statsStr,
	}
	n, _ := diff.NumDeltas()
	for i := 0; i < n; i++ {
		patch, _ := diff.Patch(i)

		s, _ := patch.String()
		r.Patches = append(r.Patches, s)

		patch.Free()
	}
	return r
}

func readBlob(repo *git.Repository, commitObj *git.Object, filepath string) (string, error) {
	c, err := commitObj.AsCommit()
	if err != nil {
		log.Print("invalid commit: ", err)
	}

	t, err := c.Tree()
	if err != nil {
		log.Print("invalid tree: ", err)
	}

	te, _ := t.EntryByPath(filepath)
	if te == nil {
		log.Print("no file: ", filepath)
		return "", fmt.Errorf("no such file/blob/tree entry %g", filepath)
	}

	f, err := repo.Lookup(te.Id)
	if err != nil {
		log.Print("invalid file: ", err)
	}

	b, err := f.AsBlob()
	if err != nil {
		log.Print("invalid blob: ", err)
	}

	return string(b.Contents()), nil
}

type Renderer func(w io.Writer, name string, i interface{})

func createPageRenderer() Renderer {
	templateFuncs := template.FuncMap{
		"humanizeTime": func(t time.Time) string {
			return humanize.Time(t)
		},
		"s_ify": func(str string, n int) string {
			if n == 1 {
				return fmt.Sprintf("%d %s", n, str)
			} else {
				return fmt.Sprintf("%d %ss", n, str)
			}
		},
		"path": func(p Pathable) string {
			return p.Path()
		},
	}

	templates := make(map[string]*template.Template)
	baseTemp := template.Must(template.ParseGlob("l/*")).Funcs(templateFuncs)

	matches, _ := filepath.Glob("t/*")
	for _, f := range matches {
		basename := filepath.Base(f)
		ext := filepath.Ext(basename)
		templates[strings.TrimSuffix(basename, ext)] = template.Must(template.Must(baseTemp.Clone()).ParseGlob(f))
	}

	return func(w io.Writer, name string, i interface{}) {
		if templates[name] == nil {
			log.Fatal("no such template ", name)
		}
		err := templates[name].ExecuteTemplate(w, "main", i)
		if err != nil {
			log.Fatal("error executing template: ", err)
		}
	}
}

func loadRepository(name string, path string) Repo {
	repo, err := git.OpenRepository(path)
	if err != nil {
		log.Fatal("failed to open repo ", path, ":", err)
	}
	desc, err := ioutil.ReadFile(path + "/.git/description")
	if err != nil {
		log.Print("failed to get repo description ", path, ":", err)
		desc = []byte("")
	}
	return Repo{
		Name:        name,
		Path:        path,
		Description: string(desc),
		_Repo:       repo,
	}
}

func getParentDir(path string) TreeEntry {
	dir, _ := filepath.Split(path)

	g := git.TreeEntry{
		Name: "..",
		Type: git.ObjectTree,
	}

	t := TreeEntry{dir, &g}
	return t
}

func getTreeEntries(t *git.Tree, treePath string) []TreeEntry {
	var r []TreeEntry
	for i := uint64(0); i < t.EntryCount(); i++ {
		r = append(r, TreeEntry{treePath, t.EntryByIndex(i)})
	}
	return r
}

// TODO: combine getTreeEntry and getSubTree into one function
func getSubTree(t *git.Tree, treePath string) []TreeEntry {
	subentry, _ := t.EntryByPath(treePath)
	if subentry.Type != git.ObjectTree {
		log.Fatal("path is not a subtree ", treePath, " - is ", subentry.Type)
	}
	subtree, _ := t.Object.Owner().LookupTree(subentry.Id)
	return append([]TreeEntry{getParentDir(treePath)}, getTreeEntries(subtree, treePath)...)
}

func getTreeEntry(t *git.Tree, treePath string) TreeEntry {
	e, _ := t.EntryByPath(treePath)
	return TreeEntry{treePath, e}
}

func (r Renderer) renderFileTree(w io.Writer, repo Repo, commit *git.Commit, path string) {
	t, _ := commit.Tree()

	readme := ""
	if buf, err := readBlob(repo._Repo, &commit.Object, "README.md"); err == nil {
		readme = string(buf)
	}

	var entries []TreeEntry
	if path == "/" || path == "" {
		entries = getTreeEntries(t, "/")
	} else {
		entries = getSubTree(t, path)
	}

	r(w, "filelist",
		struct {
			Repo    Repo
			Entries []TreeEntry
			README  string
		}{
			repo,
			entries,
			readme,
		})
}

func main() {
	repo := loadRepository("gititup", "..")
	defer repo._Repo.Free()

	render := createPageRenderer()
	masterObj, _ := repo._Repo.RevparseSingle("master")
	masterCommit, _ := masterObj.AsCommit()

	r := mux.NewRouter()
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		render.renderFileTree(w, repo, masterCommit, "/")
	})

	r.HandleFunc("/log/", func(w http.ResponseWriter, r *http.Request) {
		// TODO: figure out how to render directly to the response writer rather than returning a string
		log := getCommitLog(repo._Repo, masterObj)

		render(w, "log",
			struct {
				Repo    Repo
				Commits []Commit
			}{
				Repo:    repo,
				Commits: log,
			})
	})

	r.HandleFunc("/commit/{hash}/", func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		oid, _ := git.NewOid(vars["hash"])

		c, err := repo._Repo.LookupCommit(oid)
		if err != nil {
			log.Fatal("unable to find commit ", oid, ":", err)
		}
		defer c.Free()

		diff := getCommitDiff(repo._Repo, c)

		render(w, "diff", struct {
			Repo Repo
			Diff Diff
		}{
			repo,
			diff,
		})
	})

	r.HandleFunc("/blob/{filepath:.*}", func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		s, err := readBlob(repo._Repo, masterObj, vars["filepath"])
		if err != nil {
			s = err.Error()
		}

		render(w, "file", struct {
			Repo Repo
			Text string
		}{
			repo,
			s,
		})
	})

	r.HandleFunc("/tree/{filepath:.*}", func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		render.renderFileTree(w, repo, masterCommit, vars["filepath"])
	})

	http.Handle("/", r)
	http.ListenAndServe(":8000", nil)
}
<commit_msg>cleanup - not passing around as many random git objects as before<commit_after>package main

import (
	"fmt"
	"github.com/dustin/go-humanize"
	"github.com/gorilla/mux"
	"github.com/libgit2/git2go"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"path"
	"path/filepath"
	"strings"
	"time"
)

type TreeEntry struct {
	// extend with a pointer to the parent tree
	DirPath string

	*git.TreeEntry
}
type Pathable interface {
	Path() string
}

type Repo struct {
	Name        string
	Path        string
	Description string
	_Repo       *git.Repository
}

type Commit struct {
	_Commit *git.Commit
	Message string
	Author  string
	Hash    string
	Date    time.Time
}

type Diff struct {
	_CommitA    *git.Commit
	_CommitB    *git.Commit
	CommitAHash string
	CommitBHash string
	Stats       string
	Patches     []string
}

func (c Commit) Path() string {
	return path.Join("commit", c.Hash)
}

func (t TreeEntry) Path() string {
	if t.Type == git.ObjectBlob {
		return filepath.Join("/", "blob", t.DirPath, t.Name)
	} else if t.Type == git.ObjectTree {
		if t.DirPath == "" {
			return "/"
		} else {
			if t.Name == ".." { // TODO: simplify
				return filepath.Join("/", "tree", t.DirPath)
			} else {
				return filepath.Join("/", "tree", t.DirPath, t.Name)
			}
		}
	}
	log.Fatal("unknown type tree entry type ", t.Type)
	return ""
}

// TODO: get commit log for an arbitrary branch
func getCommitLog(repo *git.Repository, ref *git.Reference) []Commit {
	r, err := repo.Walk()
	if err != nil {
		log.Print("failed to walk repo: ", err)
	}

	r.Push(ref.Target())
	r.Sorting(git.SortTime)
	r.SimplifyFirstParent()

	id := &(git.Oid{})

	var commits []Commit
	for r.Next(id) == nil {
		g, _ := repo.LookupCommit(id)
		c := Commit{g, g.Message(), g.Committer().Name, g.Id().String(), g.Committer().When}
		commits = append(commits, c)
	}
	return commits
}

func getCommitDiff(repo *git.Repository, commit *git.Commit) Diff {
	// TODO check for multiple parent commits
	p, _ := commit.Parent(0).Tree()
	c, _ := commit.Tree()
	o, _ := git.DefaultDiffOptions()
	diff, _ := repo.DiffTreeToTree(p, c, &o)
	defer diff.Free()

	stats, _ := diff.Stats()
	statsStr, _ := stats.String(git.DiffStatsFull, 80)

	r := Diff{
		_CommitA:    commit.Parent(0),
		_CommitB:    commit,
		CommitAHash: p.Id().String(),
		CommitBHash: c.Id().String(),
		Stats:       statsStr,
	}
	n, _ := diff.NumDeltas()
	for i := 0; i < n; i++ {
		patch, _ := diff.Patch(i)

		s, _ := patch.String()
		r.Patches = append(r.Patches, s)

		patch.Free()
	}
	return r
}

func readBlob(repo *git.Repository, commit *git.Commit, filepath string) (string, error) {
	t, err := commit.Tree()
	if err != nil {
		log.Print("invalid tree: ", err)
	}

	te, _ := t.EntryByPath(filepath)
	if te == nil {
		log.Print("no file: ", filepath)
		return "", fmt.Errorf("no such file/blob/tree entry %g", filepath)
	}

	f, err := repo.Lookup(te.Id)
	if err != nil {
		log.Print("invalid file: ", err)
	}

	b, err := f.AsBlob()
	if err != nil {
		log.Print("invalid blob: ", err)
	}

	return string(b.Contents()), nil
}

type Renderer func(w io.Writer, name string, i interface{})

func createPageRenderer() Renderer {
	templateFuncs := template.FuncMap{
		"humanizeTime": func(t time.Time) string {
			return humanize.Time(t)
		},
		"s_ify": func(str string, n int) string {
			if n == 1 {
				return fmt.Sprintf("%d %s", n, str)
			} else {
				return fmt.Sprintf("%d %ss", n, str)
			}
		},
		"path": func(p Pathable) string {
			return p.Path()
		},
	}

	templates := make(map[string]*template.Template)
	baseTemp := template.Must(template.ParseGlob("l/*")).Funcs(templateFuncs)

	matches, _ := filepath.Glob("t/*")
	for _, f := range matches {
		basename := filepath.Base(f)
		ext := filepath.Ext(basename)
		templates[strings.TrimSuffix(basename, ext)] = template.Must(template.Must(baseTemp.Clone()).ParseGlob(f))
	}

	return func(w io.Writer, name string, i interface{}) {
		if templates[name] == nil {
			log.Fatal("no such template ", name)
		}
		err := templates[name].ExecuteTemplate(w, "main", i)
		if err != nil {
			log.Fatal("error executing template: ", err)
		}
	}
}

func loadRepository(name string, path string) Repo {
	repo, err := git.OpenRepository(path)
	if err != nil {
		log.Fatal("failed to open repo ", path, ":", err)
	}
	desc, err := ioutil.ReadFile(path + "/.git/description")
	if err != nil {
		log.Print("failed to get repo description ", path, ":", err)
		desc = []byte("")
	}
	return Repo{
		Name:        name,
		Path:        path,
		Description: string(desc),
		_Repo:       repo,
	}
}

func getParentDir(path string) TreeEntry {
	dir, _ := filepath.Split(path)

	g := git.TreeEntry{
		Name: "..",
		Type: git.ObjectTree,
	}

	t := TreeEntry{dir, &g}
	return t
}

func getTreeEntries(t *git.Tree, treePath string) []TreeEntry {
	var r []TreeEntry
	for i := uint64(0); i < t.EntryCount(); i++ {
		r = append(r, TreeEntry{treePath, t.EntryByIndex(i)})
	}
	return r
}

// TODO: combine getTreeEntry and getSubTree into one function
func getSubTree(t *git.Tree, treePath string) []TreeEntry {
	subentry, _ := t.EntryByPath(treePath)
	if subentry.Type != git.ObjectTree {
		log.Fatal("path is not a subtree ", treePath, " - is ", subentry.Type)
	}
	subtree, _ := t.Object.Owner().LookupTree(subentry.Id)
	return append([]TreeEntry{getParentDir(treePath)}, getTreeEntries(subtree, treePath)...)
}

func getTreeEntry(t *git.Tree, treePath string) TreeEntry {
	e, _ := t.EntryByPath(treePath)
	return TreeEntry{treePath, e}
}

func (r Renderer) renderFileTree(w io.Writer, repo Repo, commit *git.Commit, path string) {
	t, _ := commit.Tree()

	readme := ""
	if buf, err := readBlob(repo._Repo, commit, "README.md"); err == nil {
		readme = string(buf)
	}

	var entries []TreeEntry
	if path == "/" || path == "" {
		entries = getTreeEntries(t, "/")
	} else {
		entries = getSubTree(t, path)
	}

	r(w, "filelist",
		struct {
			Repo    Repo
			Entries []TreeEntry
			README  string
		}{
			repo,
			entries,
			readme,
		})
}

func main() {
	repo := loadRepository("gititup", "..")
	defer repo._Repo.Free()

	render := createPageRenderer()
	master, _ := repo._Repo.LookupBranch("master", git.BranchAll)
	firstCommitObj, _ := master.Reference.Peel(git.ObjectCommit)
	firstCommit, _ := firstCommitObj.AsCommit()

	r := mux.NewRouter()
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		render.renderFileTree(w, repo, firstCommit, "/")
	})

	r.HandleFunc("/log/", func(w http.ResponseWriter, r *http.Request) {
		log := getCommitLog(repo._Repo, master.Reference)

		render(w, "log",
			struct {
				Repo    Repo
				Commits []Commit
			}{
				Repo:    repo,
				Commits: log,
			})
	})

	r.HandleFunc("/commit/{hash}/", func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		oid, _ := git.NewOid(vars["hash"])

		c, err := repo._Repo.LookupCommit(oid)
		if err != nil {
			log.Fatal("unable to find commit ", oid, ":", err)
		}
		defer c.Free()

		diff := getCommitDiff(repo._Repo, c)

		render(w, "diff", struct {
			Repo Repo
			Diff Diff
		}{
			repo,
			diff,
		})
	})

	r.HandleFunc("/blob/{filepath:.*}", func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		s, err := readBlob(repo._Repo, firstCommit, vars["filepath"])
		if err != nil {
			s = err.Error()
		}

		render(w, "file", struct {
			Repo Repo
			Text string
		}{
			repo,
			s,
		})
	})

	r.HandleFunc("/tree/{filepath:.*}", func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		render.renderFileTree(w, repo, firstCommit, vars["filepath"])
	})

	http.Handle("/", r)
	http.ListenAndServe(":8000", nil)
}
<|endoftext|>


<commit_before>package modules

import (
	"bytes"
	"errors"
	"net"
	"strconv"
	"strings"

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

// MaxEncodedNetAddressLength is the maximum length of a NetAddress encoded
// with the encode package. 266 was chosen because the maximum length for the
// hostname is 254 + 1 for the separating colon + 5 for the port + 8 byte
// string length prefix.
const MaxEncodedNetAddressLength = 266

// A NetAddress contains the information needed to contact a peer.
type NetAddress string

// Host removes the port from a NetAddress, returning just the host. If the
// address is not of the form "host:port" the empty string is returned. The
// port will still be returned for invalid NetAddresses (e.g. "unqualified:0"
// will return "unqualified"), but in general you should only call Host on
// valid addresses.
func (na NetAddress) Host() string {
	host, _, err := net.SplitHostPort(string(na))
	// 'host' is not always the empty string if an error is returned.
	if err != nil {
		return ""
	}
	return host
}

// Port returns the NetAddress object's port number. If the address is not of
// the form "host:port" the empty string is returned. The port will still be
// returned for invalid NetAddresses (e.g. "localhost:0" will return "0"), but
// in general you should only call Port on valid addresses.
func (na NetAddress) Port() string {
	_, port, err := net.SplitHostPort(string(na))
	// 'port' will not always be the empty string if an error is returned.
	if err != nil {
		return ""
	}
	return port
}

// IsLoopback returns true for IP addresses that are on the same machine.
func (na NetAddress) IsLoopback() bool {
	host, _, err := net.SplitHostPort(string(na))
	if err != nil {
		return false
	}
	if host == "localhost" {
		return true
	}
	if ip := net.ParseIP(host); ip != nil && ip.IsLoopback() {
		return true
	}
	return false
}

// IsLocal returns true if the input IP address belongs to a local address
// range such as 192.168.x.x or 127.x.x.x
func (na NetAddress) IsLocal() bool {
	// Loopback counts as private.
	if na.IsLoopback() {
		return true
	}

	// Grab the IP address of the net address. If there is an error parsing,
	// return false, as it's not a private ip address range.
	ip := net.ParseIP(na.Host())
	if ip == nil {
		return false
	}
	ip16 := ip.To16()

	// Get the ranges of the private IP addresses.
	range1Low := net.ParseIP("10.0.0.0").To16()
	range1High := net.ParseIP("10.255.255.255").To16()
	range2Low := net.ParseIP("172.16.0.0").To16()
	range2High := net.ParseIP("172.31.255.255").To16()
	range3Low := net.ParseIP("192.168.0.0").To16()
	range3High := net.ParseIP("192.168.255.255").To16()
	range4Low := net.ParseIP("fd00:0000:0000:0000:0000:0000:0000:0000")
	range4High := net.ParseIP("fdff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")

	// Sanity check - all values should be non-nil.
	if ip16 == nil ||
		range1Low == nil || range1High == nil ||
		range2Low == nil || range2High == nil ||
		range3Low == nil || range3High == nil ||
		range4Low == nil || range4High == nil {
		panic("invalid range")
	}

	// Return true if ip16 falls between any of the above defined ranges.
	if bytes.Compare(range1Low, ip16) <= 0 && bytes.Compare(range1High, ip16) <= 0 {
		return true
	}
	if bytes.Compare(range2Low, ip16) <= 0 && bytes.Compare(range2High, ip16) <= 0 {
		return true
	}
	if bytes.Compare(range3Low, ip16) <= 0 && bytes.Compare(range3High, ip16) <= 0 {
		return true
	}
	if bytes.Compare(range4Low, ip16) <= 0 && bytes.Compare(range4High, ip16) <= 0 {
		return true
	}
	return false
}

// IsValid returns an error if the NetAddress is invalid. A valid NetAddress
// is of the form "host:port", such that "host" is either a valid IPv4/IPv6
// address or a valid hostname, and "port" is an integer in the range
// [1,65535]. Furthermore, "host" may not be a loopback address (except during
// testing). Valid IPv4 addresses, IPv6 addresses, and hostnames are detailed
// in RFCs 791, 2460, and 952, respectively.
func (na NetAddress) IsValid() error {
	host, port, err := net.SplitHostPort(string(na))
	if err != nil {
		return err
	}

	portInt, err := strconv.Atoi(port)
	if err != nil {
		return errors.New("port is not an integer")
	} else if portInt < 1 || portInt > 65535 {
		return errors.New("port is invalid")
	}

	// This check must come after the valid port check so that a host such as
	// "localhost:badport" will fail.
	if na.IsLoopback() {
		if build.Release == "testing" {
			return nil
		}
		return errors.New("host is a loopback address")
	}

	// First try to parse host as an IP address; if that fails, assume it is a
	// hostname.
	if ip := net.ParseIP(host); ip != nil {
		if ip.IsUnspecified() {
			return errors.New("host is the unspecified address")
		}
	} else {
		// Hostnames can have a trailing dot (which indicates that the hostname is
		// fully qualified), but we ignore it for validation purposes.
		if strings.HasSuffix(host, ".") {
			host = host[:len(host)-1]
		}
		if len(host) < 1 || len(host) > 253 {
			return errors.New("invalid hostname length")
		}
		labels := strings.Split(host, ".")
		if len(labels) == 1 {
			return errors.New("unqualified hostname")
		}
		for _, label := range labels {
			if len(label) < 1 || len(label) > 63 {
				return errors.New("hostname contains label with invalid length")
			}
			if strings.HasPrefix(label, "-") || strings.HasSuffix(label, "-") {
				return errors.New("hostname contains label that starts or ends with a hyphen")
			}
			for _, r := range strings.ToLower(label) {
				isLetter := 'a' <= r && r <= 'z'
				isNumber := '0' <= r && r <= '9'
				isHyphen := r == '-'
				if !(isLetter || isNumber || isHyphen) {
					return errors.New("host contains invalid characters")
				}
			}
		}
	}

	return nil
}
<commit_msg>add IsStdValid to modules.NetAddress type<commit_after>package modules

import (
	"bytes"
	"errors"
	"net"
	"strconv"
	"strings"

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

// MaxEncodedNetAddressLength is the maximum length of a NetAddress encoded
// with the encode package. 266 was chosen because the maximum length for the
// hostname is 254 + 1 for the separating colon + 5 for the port + 8 byte
// string length prefix.
const MaxEncodedNetAddressLength = 266

// A NetAddress contains the information needed to contact a peer.
type NetAddress string

// Host removes the port from a NetAddress, returning just the host. If the
// address is not of the form "host:port" the empty string is returned. The
// port will still be returned for invalid NetAddresses (e.g. "unqualified:0"
// will return "unqualified"), but in general you should only call Host on
// valid addresses.
func (na NetAddress) Host() string {
	host, _, err := net.SplitHostPort(string(na))
	// 'host' is not always the empty string if an error is returned.
	if err != nil {
		return ""
	}
	return host
}

// Port returns the NetAddress object's port number. If the address is not of
// the form "host:port" the empty string is returned. The port will still be
// returned for invalid NetAddresses (e.g. "localhost:0" will return "0"), but
// in general you should only call Port on valid addresses.
func (na NetAddress) Port() string {
	_, port, err := net.SplitHostPort(string(na))
	// 'port' will not always be the empty string if an error is returned.
	if err != nil {
		return ""
	}
	return port
}

// IsLoopback returns true for IP addresses that are on the same machine.
func (na NetAddress) IsLoopback() bool {
	host, _, err := net.SplitHostPort(string(na))
	if err != nil {
		return false
	}
	if host == "localhost" {
		return true
	}
	if ip := net.ParseIP(host); ip != nil && ip.IsLoopback() {
		return true
	}
	return false
}

// IsLocal returns true if the input IP address belongs to a local address
// range such as 192.168.x.x or 127.x.x.x
func (na NetAddress) IsLocal() bool {
	// Loopback counts as private.
	if na.IsLoopback() {
		return true
	}

	// Grab the IP address of the net address. If there is an error parsing,
	// return false, as it's not a private ip address range.
	ip := net.ParseIP(na.Host())
	if ip == nil {
		return false
	}
	ip16 := ip.To16()

	// Get the ranges of the private IP addresses.
	range1Low := net.ParseIP("10.0.0.0").To16()
	range1High := net.ParseIP("10.255.255.255").To16()
	range2Low := net.ParseIP("172.16.0.0").To16()
	range2High := net.ParseIP("172.31.255.255").To16()
	range3Low := net.ParseIP("192.168.0.0").To16()
	range3High := net.ParseIP("192.168.255.255").To16()
	range4Low := net.ParseIP("fd00:0000:0000:0000:0000:0000:0000:0000")
	range4High := net.ParseIP("fdff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")

	// Sanity check - all values should be non-nil.
	if ip16 == nil ||
		range1Low == nil || range1High == nil ||
		range2Low == nil || range2High == nil ||
		range3Low == nil || range3High == nil ||
		range4Low == nil || range4High == nil {
		panic("invalid range")
	}

	// Return true if ip16 falls between any of the above defined ranges.
	if bytes.Compare(range1Low, ip16) <= 0 && bytes.Compare(range1High, ip16) <= 0 {
		return true
	}
	if bytes.Compare(range2Low, ip16) <= 0 && bytes.Compare(range2High, ip16) <= 0 {
		return true
	}
	if bytes.Compare(range3Low, ip16) <= 0 && bytes.Compare(range3High, ip16) <= 0 {
		return true
	}
	if bytes.Compare(range4Low, ip16) <= 0 && bytes.Compare(range4High, ip16) <= 0 {
		return true
	}
	return false
}

// IsValid is an extension to IsStdValid that also forbids the loopback
// address. IsValid is being phased out in favor of allowing the loopback
// address but verifying through other means that the connection is not to
// yourself (which is the original reason that the loopback address was
// banned).
func (na NetAddress) IsValid() error {
	// Check the loopback address.
	if na.IsLoopback() && build.Release != "testing" {
		return errors.New("host is a loopback address")
	}
	return na.IsStdValid()
}

// IsStdValid returns an error if the NetAddress is invalid. A valid NetAddress
// is of the form "host:port", such that "host" is either a valid IPv4/IPv6
// address or a valid hostname, and "port" is an integer in the range
// [1,65535]. Valid IPv4 addresses, IPv6 addresses, and hostnames are detailed
// in RFCs 791, 2460, and 952, respectively.
func (na NetAddress) IsStdValid() error {
	// Verify the port number.
	host, port, err := net.SplitHostPort(string(na))
	if err != nil {
		return err
	}
	portInt, err := strconv.Atoi(port)
	if err != nil {
		return errors.New("port is not an integer")
	} else if portInt < 1 || portInt > 65535 {
		return errors.New("port is invalid")
	}

	// Loopback addresses don't always pass the requirements below, and
	// therefore must be checked separately.
	if na.IsLoopback() {
		return nil
	}

	// First try to parse host as an IP address; if that fails, assume it is a
	// hostname.
	if ip := net.ParseIP(host); ip != nil {
		if ip.IsUnspecified() {
			return errors.New("host is the unspecified address")
		}
	} else {
		// Hostnames can have a trailing dot (which indicates that the hostname is
		// fully qualified), but we ignore it for validation purposes.
		if strings.HasSuffix(host, ".") {
			host = host[:len(host)-1]
		}
		if len(host) < 1 || len(host) > 253 {
			return errors.New("invalid hostname length")
		}
		labels := strings.Split(host, ".")
		if len(labels) == 1 {
			return errors.New("unqualified hostname")
		}
		for _, label := range labels {
			if len(label) < 1 || len(label) > 63 {
				return errors.New("hostname contains label with invalid length")
			}
			if strings.HasPrefix(label, "-") || strings.HasSuffix(label, "-") {
				return errors.New("hostname contains label that starts or ends with a hyphen")
			}
			for _, r := range strings.ToLower(label) {
				isLetter := 'a' <= r && r <= 'z'
				isNumber := '0' <= r && r <= '9'
				isHyphen := r == '-'
				if !(isLetter || isNumber || isHyphen) {
					return errors.New("host contains invalid characters")
				}
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"eaciit/gdrj/model"
	"eaciit/gdrj/modules"
	"os"

	"strings"
	"time"

	"github.com/eaciit/dbox"
	"github.com/eaciit/orm/v1"
	"github.com/eaciit/toolkit"
)

var conn dbox.IConnection
var count int
var ratioTableName string

var (
	yearttxt                    = "2016"
	sourcetablename             = "salespls-summary-" + yearttxt + "-vdistrd"
	calctablename               = "salespls-summary"
	desttablename               = "salespls-summary-afterrdcopy"
	t0                          time.Time
	fiscalyear, iscount, scount int
	data                        map[string]float64
	masters                     = toolkit.M{}
)

type plalloc struct {
	Key      string
	Ref1     float64
	Current  float64
	Expect   float64
	Absorbed float64
}

type allocmap map[string]*plalloc

var plallocs = allocmap{}
var totals = allocmap{}
var ples = allocmap{}
var f = dbox.Eq("key.trxsrc", "RECLASSPROMOSPGRDMT")
var fsalesrd = dbox.Eq("key.customer_reportchannel", "RD")

func adjustAllocs(allocsmap *allocmap, key string, current, expect, absorbed, ref1 float64) {
	allocs := *allocsmap
	alloc := allocs[key]
	if alloc == nil {
		alloc = new(plalloc)
		alloc.Key = key
	}
	alloc.Current += current
	alloc.Expect += expect
	alloc.Ref1 += ref1
	alloc.Absorbed += absorbed
	allocs[key] = alloc
	*allocsmap = allocs
}

func main() {
	t0 = time.Now()

	setinitialconnection()
	defer gdrj.CloseDb()
	prepmastercalc()

	toolkit.Println("Start data query...")
	tablenames := []string{
		"salespls-summary"}

	for _, tn := range tablenames {
		e := buildRatio(tn)
		if e != nil {
			toolkit.Printfn("Build ratio error: %s - %s", tn, e.Error())
			return
		}

		e = processTable(tn)
		if e != nil {
			toolkit.Printfn("Process table error: %s - %s", tn, e.Error())
			return
		}
	}
}

func buildRatio(tn string) error {
	fiscal1 := toolkit.ToInt(yearttxt, toolkit.RoundingAuto)
	fiscal0 := fiscal1 - 1
	fiscaltxt := toolkit.Sprintf("%d-%d", fiscal0, fiscal1)
	cursor, _ := conn.NewQuery().From(calctablename).
		Where(dbox.Eq("key.date_fiscal", fiscaltxt)).
		//Group("key.customer_reportchannel").
		//Aggr(dbox.AggrSum, "PL8A", "PL8A").
		//Select().
		Cursor(nil)
	defer cursor.Close()

	i := 0
	count := cursor.Count()
	t0 := time.Now()
	mstone := 0
	total := float64(0)
	for {
		mtgtratio := toolkit.M{}
		efetch := cursor.Fetch(&mtgtratio, 1, false)
		if efetch != nil {
			break
		}
		i++
		makeProgressLog("MT/GT Ratio", i, count, 5, &mstone, t0)
		key := mtgtratio.Get("key", toolkit.M{}).(toolkit.M)
		reportchannel := key.GetString("customer_reportchannel")
		if reportchannel == "MT" || reportchannel == "GT" {
			sales := mtgtratio.GetFloat64("PL8A")
			total += sales
			adjustAllocs(&totals, reportchannel, sales, 0, 0, 0)
		}
	}
	for _, alloc := range totals {
		alloc.Expect = alloc.Current / total
		alloc.Ref1 = total
	}
	toolkit.Printfn("MT/GT Ratio: %s", toolkit.JsonString(totals))

	return nil
}

func makeProgressLog(reference string, i, count, step int, current *int, tstart time.Time) int {
	perstep := count * step / 100
	icurrent := *current
	if icurrent == 0 {
		icurrent = perstep
	}
	pct := i * 100 / count
	if i >= icurrent {
		toolkit.Printfn("Processing %s, %d of %d [%d pct] in %s",
			reference, i, count, pct, time.Since(tstart).String())
		icurrent += perstep
	}
	*current = icurrent
	return icurrent
}

func processTable(tn string) error {
	toolkit.Printfn("Start processing allocation")
	cursor, _ := conn.NewQuery().From(sourcetablename).
		Select().Cursor(nil)
	defer cursor.Close()

	//plmodels := masters["plmodel"].(map[string]*gdrj.PLModel)
	qsave := conn.NewQuery().SetConfig("multiexec", true).From(desttablename).Save()

	count := cursor.Count()
	i := 0
	step := count / 20
	mstone := step
	t0 = time.Now()
	for {
		mr := toolkit.M{}
		ef := cursor.Fetch(&mr, 1, false)
		if ef != nil {
			break
		}

		i++
		makeProgressLog("Processing", i, count, 5, &mstone, t0)
		//key := mr.Get("key", toolkit.M{}).(toolkit.M)
		//fiscal := key.GetString("date_fiscal")

		mrid := mr.GetString("_id")
		key := mr.Get("key", toolkit.M{}).(toolkit.M)
		salesvalue := mr.GetFloat64("PL8A")
		grossvalue := mr.GetFloat64("PL2")

		for channel, total := range totals {
			mrk := toolkit.M{}
			mrkkey := key
			mrk.Set("key", mrkkey)
			mrkkey.Set("trxsrc", "pushrdreversesbymks")
			mrkkey.Set("customer_reportchannel", channel)
			mrkkey.Set("customer_channelname", channel)
			if channel == "MT" {
				mrk.Set("customer_channelid", "I3")
			} else if channel == "GT" {
				mrk.Set("customer_channelid", "I2")
			}

			mrsales := -salesvalue * total.Expect
			mrgross := -grossvalue * total.Expect
			mrdiscount := mrsales - mrgross
			mrk.Set("PL1", mrgross)
			mrk.Set("PL7", mrdiscount)
			mrk.Set("PL8A", mrsales)

			mrk.Set("_id", toolkit.Sprintf("%s|pushrdreverse|%s", mrid, channel))
			gdrj.CalcSum(mrk, masters)
			esavereverse := qsave.Exec(toolkit.M{}.Set("data", mrk))
			if esavereverse != nil {
				return esavereverse
			}
		}

		for k, v := range mr {
			if strings.HasPrefix(k, "PL") {
				if k == "PL8" || k == "PL2" {
					mr.Set(k, v)
				} else {
					mr.Set(k, float64(0))
				}
			}
		}

		key.Set("trxsrc", "rdsbymks")
		mr.Set("key", key)
		gdrj.CalcSum(mr, masters)
		esave := qsave.Exec(toolkit.M{}.Set("data", mr))
		if esave != nil {
			return esave
		}
	}
	return nil
}

func buildmap(holder interface{},
	fnModel func() orm.IModel,
	filter *dbox.Filter,
	fnIter func(holder interface{}, obj interface{})) interface{} {
	crx, ecrx := gdrj.Find(fnModel(), filter, nil)
	if ecrx != nil {
		toolkit.Printfn("Cursor Error: %s", ecrx.Error())
		os.Exit(100)
	}
	defer crx.Close()
	for {
		s := fnModel()
		e := crx.Fetch(s, 1, false)
		if e != nil {
			break
		}
		fnIter(holder, s)
	}
	return holder
}

func prepmastercalc() {
	toolkit.Println("--> PL MODEL")
	masters.Set("plmodel", buildmap(map[string]*gdrj.PLModel{},
		func() orm.IModel {
			return new(gdrj.PLModel)
		},
		nil,
		func(holder, obj interface{}) {
			h := holder.(map[string]*gdrj.PLModel)
			o := obj.(*gdrj.PLModel)
			h[o.ID] = o
		}).(map[string]*gdrj.PLModel))
}

func setinitialconnection() {
	var err error
	conn, err = modules.GetDboxIConnection("db_godrej")

	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}

	err = gdrj.SetDb(conn)
	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}
}
<commit_msg>upd k<commit_after>package main

import (
	"eaciit/gdrj/model"
	"eaciit/gdrj/modules"
	"os"

	"strings"
	"time"

	"github.com/eaciit/dbox"
	"github.com/eaciit/orm/v1"
	"github.com/eaciit/toolkit"
)

var conn dbox.IConnection
var count int
var ratioTableName string

var (
	yearttxt                    = "2016"
	sourcetablename             = "salespls-summary-" + yearttxt + "-vdistrd"
	calctablename               = "salespls-summary"
	desttablename               = "salespls-summary-afterrdcopy"
	t0                          time.Time
	fiscalyear, iscount, scount int
	data                        map[string]float64
	masters                     = toolkit.M{}
)

type plalloc struct {
	Key      string
	Ref1     float64
	Current  float64
	Expect   float64
	Absorbed float64
}

type allocmap map[string]*plalloc

var plallocs = allocmap{}
var totals = allocmap{}
var ples = allocmap{}
var f = dbox.Eq("key.trxsrc", "RECLASSPROMOSPGRDMT")
var fsalesrd = dbox.Eq("key.customer_reportchannel", "RD")

func adjustAllocs(allocsmap *allocmap, key string, current, expect, absorbed, ref1 float64) {
	allocs := *allocsmap
	alloc := allocs[key]
	if alloc == nil {
		alloc = new(plalloc)
		alloc.Key = key
	}
	alloc.Current += current
	alloc.Expect += expect
	alloc.Ref1 += ref1
	alloc.Absorbed += absorbed
	allocs[key] = alloc
	*allocsmap = allocs
}

func main() {
	t0 = time.Now()

	setinitialconnection()
	defer gdrj.CloseDb()
	prepmastercalc()

	toolkit.Println("Start data query...")
	tablenames := []string{
		"salespls-summary"}

	for _, tn := range tablenames {
		e := buildRatio(tn)
		if e != nil {
			toolkit.Printfn("Build ratio error: %s - %s", tn, e.Error())
			return
		}

		e = processTable(tn)
		if e != nil {
			toolkit.Printfn("Process table error: %s - %s", tn, e.Error())
			return
		}
	}
}

func buildRatio(tn string) error {
	fiscal1 := toolkit.ToInt(yearttxt, toolkit.RoundingAuto)
	fiscal0 := fiscal1 - 1
	fiscaltxt := toolkit.Sprintf("%d-%d", fiscal0, fiscal1)
	cursor, _ := conn.NewQuery().From(calctablename).
		Where(dbox.Eq("key.date_fiscal", fiscaltxt)).
		//Group("key.customer_reportchannel").
		//Aggr(dbox.AggrSum, "PL8A", "PL8A").
		//Select().
		Cursor(nil)
	defer cursor.Close()

	i := 0
	count := cursor.Count()
	t0 := time.Now()
	mstone := 0
	total := float64(0)
	for {
		mtgtratio := toolkit.M{}
		efetch := cursor.Fetch(&mtgtratio, 1, false)
		if efetch != nil {
			break
		}
		i++
		makeProgressLog("MT/GT Ratio", i, count, 5, &mstone, t0)
		key := mtgtratio.Get("key", toolkit.M{}).(toolkit.M)
		reportchannel := key.GetString("customer_reportchannel")
		if reportchannel == "MT" || reportchannel == "GT" {
			sales := mtgtratio.GetFloat64("PL8A")
			total += sales
			adjustAllocs(&totals, reportchannel, sales, 0, 0, 0)
		}
	}
	for _, alloc := range totals {
		alloc.Expect = alloc.Current / total
		alloc.Ref1 = total
	}
	toolkit.Printfn("MT/GT Ratio: %s", toolkit.JsonString(totals))

	return nil
}

func makeProgressLog(reference string, i, count, step int, current *int, tstart time.Time) int {
	perstep := count * step / 100
	icurrent := *current
	if icurrent == 0 {
		icurrent = perstep
	}
	pct := i * 100 / count
	if i >= icurrent {
		toolkit.Printfn("Processing %s, %d of %d [%d pct] in %s",
			reference, i, count, pct, time.Since(tstart).String())
		icurrent += perstep
	}
	*current = icurrent
	return icurrent
}

func processTable(tn string) error {
	toolkit.Printfn("Start processing allocation")
	cursor, _ := conn.NewQuery().From(sourcetablename).
		Select().Cursor(nil)
	defer cursor.Close()

	//plmodels := masters["plmodel"].(map[string]*gdrj.PLModel)
	qsave := conn.NewQuery().SetConfig("multiexec", true).From(desttablename).Save()

	count := cursor.Count()
	i := 0
	step := count / 20
	mstone := step
	t0 = time.Now()
	for {
		mr := toolkit.M{}
		ef := cursor.Fetch(&mr, 1, false)
		if ef != nil {
			break
		}

		i++
		makeProgressLog("Processing", i, count, 5, &mstone, t0)
		//key := mr.Get("key", toolkit.M{}).(toolkit.M)
		//fiscal := key.GetString("date_fiscal")

		mrid := mr.GetString("_id")
		key := mr.Get("key", toolkit.M{}).(toolkit.M)
		salesvalue := mr.GetFloat64("PL8A")
		grossvalue := mr.GetFloat64("PL2")

		for channel, total := range totals {
			mrk := toolkit.M{}
			mrkkey := toolkit.M{}
			for k, v := range key {
				mrkkey.Set(k, v)
			}
			mrkkey.Set("trxsrc", "pushrdreversesbymks")
			mrkkey.Set("customer_reportchannel", channel)
			mrkkey.Set("customer_channelname", channel)
			if channel == "MT" {
				mrk.Set("customer_channelid", "I3")
			} else if channel == "GT" {
				mrk.Set("customer_channelid", "I2")
			}

			mrsales := -salesvalue * total.Expect
			mrgross := -grossvalue * total.Expect
			mrdiscount := mrsales - mrgross
			mrk.Set("key", mrkkey)
			mrk.Set("PL1", mrgross)
			mrk.Set("PL7", mrdiscount)
			mrk.Set("PL8A", mrsales)

			mrk.Set("_id", toolkit.Sprintf("%s|pushrdreverse|%s", mrid, channel))
			gdrj.CalcSum(mrk, masters)
			esavereverse := qsave.Exec(toolkit.M{}.Set("data", mrk))
			if esavereverse != nil {
				return esavereverse
			}
		}

		for k, v := range mr {
			if strings.HasPrefix(k, "PL") {
				if k == "PL8" || k == "PL2" {
					mr.Set(k, v)
				} else {
					mr.Set(k, float64(0))
				}
			}
		}

		key.Set("trxsrc", "rdsbymks")
		mr.Set("key", key)
		gdrj.CalcSum(mr, masters)
		esave := qsave.Exec(toolkit.M{}.Set("data", mr))
		if esave != nil {
			return esave
		}
	}
	return nil
}

func buildmap(holder interface{},
	fnModel func() orm.IModel,
	filter *dbox.Filter,
	fnIter func(holder interface{}, obj interface{})) interface{} {
	crx, ecrx := gdrj.Find(fnModel(), filter, nil)
	if ecrx != nil {
		toolkit.Printfn("Cursor Error: %s", ecrx.Error())
		os.Exit(100)
	}
	defer crx.Close()
	for {
		s := fnModel()
		e := crx.Fetch(s, 1, false)
		if e != nil {
			break
		}
		fnIter(holder, s)
	}
	return holder
}

func prepmastercalc() {
	toolkit.Println("--> PL MODEL")
	masters.Set("plmodel", buildmap(map[string]*gdrj.PLModel{},
		func() orm.IModel {
			return new(gdrj.PLModel)
		},
		nil,
		func(holder, obj interface{}) {
			h := holder.(map[string]*gdrj.PLModel)
			o := obj.(*gdrj.PLModel)
			h[o.ID] = o
		}).(map[string]*gdrj.PLModel))
}

func setinitialconnection() {
	var err error
	conn, err = modules.GetDboxIConnection("db_godrej")

	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}

	err = gdrj.SetDb(conn)
	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}
}
<|endoftext|>


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

package sessions

import (
	"bytes"
	"encoding/gob"
	"net/http"
	"net/http/httptest"
	"testing"
)

// NewRecorder returns an initialized ResponseRecorder.
func NewRecorder() *httptest.ResponseRecorder {
	return &httptest.ResponseRecorder{
		HeaderMap: make(http.Header),
		Body:      new(bytes.Buffer),
	}
}

// DefaultRemoteAddr is the default remote address to return in RemoteAddr if
// an explicit DefaultRemoteAddr isn't set on ResponseRecorder.
const DefaultRemoteAddr = "1.2.3.4"

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

type FlashMessage struct {
	Type    int
	Message string
}

func TestFlashes(t *testing.T) {
	var req *http.Request
	var rsp *httptest.ResponseRecorder
	var hdr http.Header
	var err error
	var ok bool
	var cookies []string
	var session *Session
	var flashes []interface{}

	store := NewCookieStore([]byte("secret-key"))

	// Round 1 ----------------------------------------------------------------

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
	rsp = NewRecorder()
	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}
	// Get a flash.
	flashes = session.Flashes()
	if len(flashes) != 0 {
		t.Errorf("Expected empty flashes; Got %v", flashes)
	}
	// Add some flashes.
	session.AddFlash("foo")
	session.AddFlash("bar")
	// Custom key.
	session.AddFlash("baz", "custom_key")
	// Save.
	if err = Save(req, rsp); err != nil {
		t.Fatalf("Error saving session: %v", err)
	}
	hdr = rsp.Header()
	cookies, ok = hdr["Set-Cookie"]
	if !ok || len(cookies) != 1 {
		t.Fatal("No cookies. Header:", hdr)
	}

	if _, err = store.Get(req, "session:key"); err.Error() != "sessions: invalid character in cookie name: session:key" {
		t.Fatalf("Expected error due to invalid cookie name")
	}

	// Round 2 ----------------------------------------------------------------

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
	req.Header.Add("Cookie", cookies[0])
	rsp = NewRecorder()
	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}
	// Check all saved values.
	flashes = session.Flashes()
	if len(flashes) != 2 {
		t.Fatalf("Expected flashes; Got %v", flashes)
	}
	if flashes[0] != "foo" || flashes[1] != "bar" {
		t.Errorf("Expected foo,bar; Got %v", flashes)
	}
	flashes = session.Flashes()
	if len(flashes) != 0 {
		t.Errorf("Expected dumped flashes; Got %v", flashes)
	}
	// Custom key.
	flashes = session.Flashes("custom_key")
	if len(flashes) != 1 {
		t.Errorf("Expected flashes; Got %v", flashes)
	} else if flashes[0] != "baz" {
		t.Errorf("Expected baz; Got %v", flashes)
	}
	flashes = session.Flashes("custom_key")
	if len(flashes) != 0 {
		t.Errorf("Expected dumped flashes; Got %v", flashes)
	}

	// Round 3 ----------------------------------------------------------------
	// Custom type

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
	rsp = NewRecorder()
	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}
	// Get a flash.
	flashes = session.Flashes()
	if len(flashes) != 0 {
		t.Errorf("Expected empty flashes; Got %v", flashes)
	}
	// Add some flashes.
	session.AddFlash(&FlashMessage{42, "foo"})
	// Save.
	if err = Save(req, rsp); err != nil {
		t.Fatalf("Error saving session: %v", err)
	}
	hdr = rsp.Header()
	cookies, ok = hdr["Set-Cookie"]
	if !ok || len(cookies) != 1 {
		t.Fatal("No cookies. Header:", hdr)
	}

	// Round 4 ----------------------------------------------------------------
	// Custom type

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
	req.Header.Add("Cookie", cookies[0])
	rsp = NewRecorder()
	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}
	// Check all saved values.
	flashes = session.Flashes()
	if len(flashes) != 1 {
		t.Fatalf("Expected flashes; Got %v", flashes)
	}
	custom := flashes[0].(FlashMessage)
	if custom.Type != 42 || custom.Message != "foo" {
		t.Errorf("Expected %#v, got %#v", FlashMessage{42, "foo"}, custom)
	}

	// Round 5 ----------------------------------------------------------------
	// Check if a request shallow copy resets the request context data store.

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)

	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}

	// Put a test value into the session data store.
	session.Values["test"] = "test-value"

	// Create a shallow copy of the request.
	req = req.WithContext(req.Context())

	// Get the session again.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}

	// Check if the previous inserted value still exists.
	if session.Values["test"] == nil {
		t.Fatalf("Session test value is lost in the request context!")
	}

	// Check if the previous inserted value has the same value.
	if session.Values["test"] != "test-value" {
		t.Fatalf("Session test value is changed in the request context!")
	}
}

func TestCookieStoreMapPanic(t *testing.T) {
	defer func() {
		err := recover()
		if err != nil {
			t.Fatal(err)
		}
	}()

	store := NewCookieStore([]byte("aaa0defe5d2839cbc46fc4f080cd7adc"))
	req, err := http.NewRequest("GET", "http://www.example.com", nil)
	if err != nil {
		t.Fatal("failed to create request", err)
	}
	w := httptest.NewRecorder()

	session := NewSession(store, "hello")

	session.Values["data"] = "hello-world"

	err = session.Save(req, w)
	if err != nil {
		t.Fatal("failed to save session", err)
	}
}

func init() {
	gob.Register(FlashMessage{})
}
<commit_msg>Removed unused global var (#199)<commit_after>// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sessions

import (
	"bytes"
	"encoding/gob"
	"net/http"
	"net/http/httptest"
	"testing"
)

// NewRecorder returns an initialized ResponseRecorder.
func NewRecorder() *httptest.ResponseRecorder {
	return &httptest.ResponseRecorder{
		HeaderMap: make(http.Header),
		Body:      new(bytes.Buffer),
	}
}

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

type FlashMessage struct {
	Type    int
	Message string
}

func TestFlashes(t *testing.T) {
	var req *http.Request
	var rsp *httptest.ResponseRecorder
	var hdr http.Header
	var err error
	var ok bool
	var cookies []string
	var session *Session
	var flashes []interface{}

	store := NewCookieStore([]byte("secret-key"))

	// Round 1 ----------------------------------------------------------------

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
	rsp = NewRecorder()
	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}
	// Get a flash.
	flashes = session.Flashes()
	if len(flashes) != 0 {
		t.Errorf("Expected empty flashes; Got %v", flashes)
	}
	// Add some flashes.
	session.AddFlash("foo")
	session.AddFlash("bar")
	// Custom key.
	session.AddFlash("baz", "custom_key")
	// Save.
	if err = Save(req, rsp); err != nil {
		t.Fatalf("Error saving session: %v", err)
	}
	hdr = rsp.Header()
	cookies, ok = hdr["Set-Cookie"]
	if !ok || len(cookies) != 1 {
		t.Fatal("No cookies. Header:", hdr)
	}

	if _, err = store.Get(req, "session:key"); err.Error() != "sessions: invalid character in cookie name: session:key" {
		t.Fatalf("Expected error due to invalid cookie name")
	}

	// Round 2 ----------------------------------------------------------------

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
	req.Header.Add("Cookie", cookies[0])
	rsp = NewRecorder()
	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}
	// Check all saved values.
	flashes = session.Flashes()
	if len(flashes) != 2 {
		t.Fatalf("Expected flashes; Got %v", flashes)
	}
	if flashes[0] != "foo" || flashes[1] != "bar" {
		t.Errorf("Expected foo,bar; Got %v", flashes)
	}
	flashes = session.Flashes()
	if len(flashes) != 0 {
		t.Errorf("Expected dumped flashes; Got %v", flashes)
	}
	// Custom key.
	flashes = session.Flashes("custom_key")
	if len(flashes) != 1 {
		t.Errorf("Expected flashes; Got %v", flashes)
	} else if flashes[0] != "baz" {
		t.Errorf("Expected baz; Got %v", flashes)
	}
	flashes = session.Flashes("custom_key")
	if len(flashes) != 0 {
		t.Errorf("Expected dumped flashes; Got %v", flashes)
	}

	// Round 3 ----------------------------------------------------------------
	// Custom type

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
	rsp = NewRecorder()
	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}
	// Get a flash.
	flashes = session.Flashes()
	if len(flashes) != 0 {
		t.Errorf("Expected empty flashes; Got %v", flashes)
	}
	// Add some flashes.
	session.AddFlash(&FlashMessage{42, "foo"})
	// Save.
	if err = Save(req, rsp); err != nil {
		t.Fatalf("Error saving session: %v", err)
	}
	hdr = rsp.Header()
	cookies, ok = hdr["Set-Cookie"]
	if !ok || len(cookies) != 1 {
		t.Fatal("No cookies. Header:", hdr)
	}

	// Round 4 ----------------------------------------------------------------
	// Custom type

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
	req.Header.Add("Cookie", cookies[0])
	rsp = NewRecorder()
	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}
	// Check all saved values.
	flashes = session.Flashes()
	if len(flashes) != 1 {
		t.Fatalf("Expected flashes; Got %v", flashes)
	}
	custom := flashes[0].(FlashMessage)
	if custom.Type != 42 || custom.Message != "foo" {
		t.Errorf("Expected %#v, got %#v", FlashMessage{42, "foo"}, custom)
	}

	// Round 5 ----------------------------------------------------------------
	// Check if a request shallow copy resets the request context data store.

	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)

	// Get a session.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}

	// Put a test value into the session data store.
	session.Values["test"] = "test-value"

	// Create a shallow copy of the request.
	req = req.WithContext(req.Context())

	// Get the session again.
	if session, err = store.Get(req, "session-key"); err != nil {
		t.Fatalf("Error getting session: %v", err)
	}

	// Check if the previous inserted value still exists.
	if session.Values["test"] == nil {
		t.Fatalf("Session test value is lost in the request context!")
	}

	// Check if the previous inserted value has the same value.
	if session.Values["test"] != "test-value" {
		t.Fatalf("Session test value is changed in the request context!")
	}
}

func TestCookieStoreMapPanic(t *testing.T) {
	defer func() {
		err := recover()
		if err != nil {
			t.Fatal(err)
		}
	}()

	store := NewCookieStore([]byte("aaa0defe5d2839cbc46fc4f080cd7adc"))
	req, err := http.NewRequest("GET", "http://www.example.com", nil)
	if err != nil {
		t.Fatal("failed to create request", err)
	}
	w := httptest.NewRecorder()

	session := NewSession(store, "hello")

	session.Values["data"] = "hello-world"

	err = session.Save(req, w)
	if err != nil {
		t.Fatal("failed to save session", err)
	}
}

func init() {
	gob.Register(FlashMessage{})
}
<|endoftext|>


<commit_before>package main

import (
	"html/template"
	"net/http"
	"strconv"
	"strings"
	"time"

	// "github.com/gin-contrib/static"
	"github.com/gin-contrib/multitemplate"
	"github.com/gin-gonic/gin"
)

func serve(port string) {
	gin.SetMode(gin.ReleaseMode)
	router := gin.Default()
	router.HTMLRender = loadTemplates("index.tmpl")
	// router.Use(static.Serve("/static/", static.LocalFile("./static", true)))
	router.GET("/", func(c *gin.Context) {
		c.Redirect(302, "/"+randomAlliterateCombo())
	})
	router.GET("/:page", func(c *gin.Context) {
		page := c.Param("page")
		c.Redirect(302, "/"+page+"/edit")
	})
	router.GET("/:page/*command", handlePageRequest)
	router.POST("/update", handlePageUpdate)
	router.POST("/prime", handlePrime)
	router.POST("/lock", handleLock)
	router.POST("/encrypt", handleEncrypt)
	router.DELETE("/listitem", deleteListItem)

	router.Run(":" + port)
}

func loadTemplates(list ...string) multitemplate.Render {
	r := multitemplate.New()

	for _, x := range list {
		templateString, err := Asset("templates/" + x)
		if err != nil {
			panic(err)
		}

		tmplMessage, err := template.New(x).Parse(string(templateString))
		if err != nil {
			panic(err)
		}

		r.Add(x, tmplMessage)
	}

	return r
}

func handlePageRequest(c *gin.Context) {
	page := c.Param("page")
	command := c.Param("command")

	// Serve static content from memory
	if page == "static" {
		filename := page + command
		data, err := Asset(filename)
		if err != nil {
			c.String(http.StatusInternalServerError, "Could not find data")
		}
		c.Data(http.StatusOK, contentType(filename), data)
		return
	}

	version := c.DefaultQuery("version", "ajksldfjl")
	p := Open(page)
	if p.IsPrimedForSelfDestruct && !p.IsLocked && !p.IsEncrypted {
		p.Update("*This page has now self-destructed.*\n\n" + p.Text.GetCurrent())
		p.Erase()
	}
	if command == "/erase" {
		if !p.IsLocked && !p.IsEncrypted {
			p.Erase()
			c.Redirect(302, "/"+page+"/edit")
			return
		} else {
			c.Redirect(302, "/"+page+"/view")
			return
		}
	}
	rawText := p.Text.GetCurrent()
	rawHTML := p.RenderedPage

	// Check to see if an old version is requested
	versionInt, versionErr := strconv.Atoi(version)
	if versionErr == nil && versionInt > 0 {
		versionText, err := p.Text.GetPreviousByTimestamp(int64(versionInt))
		if err == nil {
			rawText = versionText
			rawHTML = GithubMarkdownToHTML(rawText)
		}
	}
	versionsInt64 := p.Text.GetMajorSnapshots()
	versionsText := make([]string, len(versionsInt64))
	for i, v := range versionsInt64 {
		versionsText[i] = time.Unix(v/1000000000, 0).String()
	}

	if command == "/raw" {
		c.Writer.Header().Set("Content-Type", contentType(p.Name))
		c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
		c.Writer.Header().Set("Access-Control-Max-Age", "86400")
		c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE")
		c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, X-Max")
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.Data(200, contentType(p.Name), []byte(rawText))
		return
	}
	log.Debug(command)
	c.HTML(http.StatusOK, "index.tmpl", gin.H{
		"EditPage":     command == "/edit",
		"ViewPage":     command == "/view",
		"ListPage":     command == "/list",
		"HistoryPage":  command == "/history",
		"Page":         p.Name,
		"RenderedPage": template.HTML([]byte(rawHTML)),
		"RawPage":      rawText,
		"Versions":     versionsInt64,
		"VersionsText": versionsText,
		"IsLocked":     p.IsLocked,
		"IsEncrypted":  p.IsEncrypted,
		"ListItems":    renderList(rawText),
	})
}

func handlePageUpdate(c *gin.Context) {
	type QueryJSON struct {
		Page    string `json:"page"`
		NewText string `json:"new_text"`
	}
	var json QueryJSON
	if c.BindJSON(&json) != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Wrong JSON"})
		return
	}
	if len(json.NewText) > 100000 {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Too much"})
		return
	}
	log.Trace("Update: %v", json)
	p := Open(json.Page)
	var message string
	if p.IsLocked {
		message = "Locked"
	} else if p.IsEncrypted {
		message = "Encrypted"
	} else {
		p.Update(json.NewText)
		p.Save()
		message = "Saved"
	}
	c.JSON(http.StatusOK, gin.H{"success": false, "message": message})
}

func handlePrime(c *gin.Context) {
	type QueryJSON struct {
		Page string `json:"page"`
	}
	var json QueryJSON
	if c.BindJSON(&json) != nil {
		c.String(http.StatusBadRequest, "Problem binding keys")
		return
	}
	log.Trace("Update: %v", json)
	p := Open(json.Page)
	if p.IsLocked {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Locked"})
		return
	} else if p.IsEncrypted {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Encrypted"})
		return
	}
	p.IsPrimedForSelfDestruct = true
	p.Save()
	c.JSON(http.StatusOK, gin.H{"success": true, "message": "Primed"})
}

func handleLock(c *gin.Context) {
	type QueryJSON struct {
		Page       string `json:"page"`
		Passphrase string `json:"passphrase"`
	}

	var json QueryJSON
	if c.BindJSON(&json) != nil {
		c.String(http.StatusBadRequest, "Problem binding keys")
		return
	}
	p := Open(json.Page)
	if p.IsEncrypted {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Encrypted"})
		return
	}
	var message string
	if p.IsLocked {
		err2 := CheckPasswordHash(json.Passphrase, p.PassphraseToUnlock)
		if err2 != nil {
			c.JSON(http.StatusOK, gin.H{"success": false, "message": "Can't unlock"})
			return
		}
		p.IsLocked = false
		message = "Unlocked"
	} else {
		p.IsLocked = true
		p.PassphraseToUnlock = HashPassword(json.Passphrase)
		message = "Locked"
	}
	p.Save()
	c.JSON(http.StatusOK, gin.H{"success": true, "message": message})
}

func handleEncrypt(c *gin.Context) {
	type QueryJSON struct {
		Page       string `json:"page"`
		Passphrase string `json:"passphrase"`
	}

	var json QueryJSON
	if c.BindJSON(&json) != nil {
		c.String(http.StatusBadRequest, "Problem binding keys")
		return
	}
	p := Open(json.Page)
	if p.IsLocked {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Locked"})
		return
	}
	q := Open(json.Page)
	var message string
	if p.IsEncrypted {
		decrypted, err2 := DecryptString(p.Text.GetCurrent(), json.Passphrase)
		if err2 != nil {
			c.JSON(http.StatusOK, gin.H{"success": false, "message": "Wrong password"})
			return
		}
		q.Erase()
		q = Open(json.Page)
		q.Update(decrypted)
		q.IsEncrypted = false
		q.IsLocked = p.IsLocked
		q.IsPrimedForSelfDestruct = p.IsPrimedForSelfDestruct
		message = "Decrypted"
	} else {
		currentText := p.Text.GetCurrent()
		encrypted, _ := EncryptString(currentText, json.Passphrase)
		q.Erase()
		q = Open(json.Page)
		q.Update(encrypted)
		q.IsEncrypted = true
		q.IsLocked = p.IsLocked
		q.IsPrimedForSelfDestruct = p.IsPrimedForSelfDestruct
		message = "Encrypted"
	}
	q.Save()
	c.JSON(http.StatusOK, gin.H{"success": true, "message": message})
}

func deleteListItem(c *gin.Context) {
	lineNum, err := strconv.Atoi(c.DefaultQuery("lineNum", "None"))
	page := c.Query("page") // shortcut for c.Request.URL.Query().Get("lastname")
	if err == nil {
		p := Open(page)

		_, listItems := reorderList(p.Text.GetCurrent())
		newText := p.Text.GetCurrent()
		for i, lineString := range listItems {
			// fmt.Println(i, lineString, lineNum)
			if i+1 == lineNum {
				// fmt.Println("MATCHED")
				if strings.Contains(lineString, "~~") == false {
					// fmt.Println(p.Text, "("+lineString[2:]+"\n"+")", "~~"+lineString[2:]+"~~"+"\n")
					newText = strings.Replace(newText+"\n", lineString[2:]+"\n", "~~"+strings.TrimSpace(lineString[2:])+"~~"+"\n", 1)
				} else {
					newText = strings.Replace(newText+"\n", lineString[2:]+"\n", lineString[4:len(lineString)-2]+"\n", 1)
				}
				p.Update(newText)
				break
			}
		}

		c.JSON(200, gin.H{
			"success": true,
			"message": "Done.",
		})
	} else {
		c.JSON(200, gin.H{
			"success": false,
			"message": err.Error(),
		})
	}
}
<commit_msg>Update GetMajorSnapshots API<commit_after>package main

import (
	"html/template"
	"net/http"
	"strconv"
	"strings"
	"time"

	// "github.com/gin-contrib/static"
	"github.com/gin-contrib/multitemplate"
	"github.com/gin-gonic/gin"
)

func serve(port string) {
	gin.SetMode(gin.ReleaseMode)
	router := gin.Default()
	router.HTMLRender = loadTemplates("index.tmpl")
	// router.Use(static.Serve("/static/", static.LocalFile("./static", true)))
	router.GET("/", func(c *gin.Context) {
		c.Redirect(302, "/"+randomAlliterateCombo())
	})
	router.GET("/:page", func(c *gin.Context) {
		page := c.Param("page")
		c.Redirect(302, "/"+page+"/edit")
	})
	router.GET("/:page/*command", handlePageRequest)
	router.POST("/update", handlePageUpdate)
	router.POST("/prime", handlePrime)
	router.POST("/lock", handleLock)
	router.POST("/encrypt", handleEncrypt)
	router.DELETE("/listitem", deleteListItem)

	router.Run(":" + port)
}

func loadTemplates(list ...string) multitemplate.Render {
	r := multitemplate.New()

	for _, x := range list {
		templateString, err := Asset("templates/" + x)
		if err != nil {
			panic(err)
		}

		tmplMessage, err := template.New(x).Parse(string(templateString))
		if err != nil {
			panic(err)
		}

		r.Add(x, tmplMessage)
	}

	return r
}

func handlePageRequest(c *gin.Context) {
	page := c.Param("page")
	command := c.Param("command")

	// Serve static content from memory
	if page == "static" {
		filename := page + command
		data, err := Asset(filename)
		if err != nil {
			c.String(http.StatusInternalServerError, "Could not find data")
		}
		c.Data(http.StatusOK, contentType(filename), data)
		return
	}

	version := c.DefaultQuery("version", "ajksldfjl")
	p := Open(page)
	if p.IsPrimedForSelfDestruct && !p.IsLocked && !p.IsEncrypted {
		p.Update("*This page has now self-destructed.*\n\n" + p.Text.GetCurrent())
		p.Erase()
	}
	if command == "/erase" {
		if !p.IsLocked && !p.IsEncrypted {
			p.Erase()
			c.Redirect(302, "/"+page+"/edit")
			return
		} else {
			c.Redirect(302, "/"+page+"/view")
			return
		}
	}
	rawText := p.Text.GetCurrent()
	rawHTML := p.RenderedPage

	// Check to see if an old version is requested
	versionInt, versionErr := strconv.Atoi(version)
	if versionErr == nil && versionInt > 0 {
		versionText, err := p.Text.GetPreviousByTimestamp(int64(versionInt))
		if err == nil {
			rawText = versionText
			rawHTML = GithubMarkdownToHTML(rawText)
		}
	}
	versionsInt64 := p.Text.GetMajorSnapshots(60) // get snapshots 60 seconds apart
	versionsText := make([]string, len(versionsInt64))
	for i, v := range versionsInt64 {
		versionsText[i] = time.Unix(v/1000000000, 0).String()
	}

	if command == "/raw" {
		c.Writer.Header().Set("Content-Type", contentType(p.Name))
		c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
		c.Writer.Header().Set("Access-Control-Max-Age", "86400")
		c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE")
		c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, X-Max")
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.Data(200, contentType(p.Name), []byte(rawText))
		return
	}
	log.Debug(command)
	c.HTML(http.StatusOK, "index.tmpl", gin.H{
		"EditPage":     command == "/edit",
		"ViewPage":     command == "/view",
		"ListPage":     command == "/list",
		"HistoryPage":  command == "/history",
		"Page":         p.Name,
		"RenderedPage": template.HTML([]byte(rawHTML)),
		"RawPage":      rawText,
		"Versions":     versionsInt64,
		"VersionsText": versionsText,
		"IsLocked":     p.IsLocked,
		"IsEncrypted":  p.IsEncrypted,
		"ListItems":    renderList(rawText),
	})
}

func handlePageUpdate(c *gin.Context) {
	type QueryJSON struct {
		Page    string `json:"page"`
		NewText string `json:"new_text"`
	}
	var json QueryJSON
	if c.BindJSON(&json) != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Wrong JSON"})
		return
	}
	if len(json.NewText) > 100000 {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Too much"})
		return
	}
	log.Trace("Update: %v", json)
	p := Open(json.Page)
	var message string
	if p.IsLocked {
		message = "Locked"
	} else if p.IsEncrypted {
		message = "Encrypted"
	} else {
		p.Update(json.NewText)
		p.Save()
		message = "Saved"
	}
	c.JSON(http.StatusOK, gin.H{"success": false, "message": message})
}

func handlePrime(c *gin.Context) {
	type QueryJSON struct {
		Page string `json:"page"`
	}
	var json QueryJSON
	if c.BindJSON(&json) != nil {
		c.String(http.StatusBadRequest, "Problem binding keys")
		return
	}
	log.Trace("Update: %v", json)
	p := Open(json.Page)
	if p.IsLocked {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Locked"})
		return
	} else if p.IsEncrypted {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Encrypted"})
		return
	}
	p.IsPrimedForSelfDestruct = true
	p.Save()
	c.JSON(http.StatusOK, gin.H{"success": true, "message": "Primed"})
}

func handleLock(c *gin.Context) {
	type QueryJSON struct {
		Page       string `json:"page"`
		Passphrase string `json:"passphrase"`
	}

	var json QueryJSON
	if c.BindJSON(&json) != nil {
		c.String(http.StatusBadRequest, "Problem binding keys")
		return
	}
	p := Open(json.Page)
	if p.IsEncrypted {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Encrypted"})
		return
	}
	var message string
	if p.IsLocked {
		err2 := CheckPasswordHash(json.Passphrase, p.PassphraseToUnlock)
		if err2 != nil {
			c.JSON(http.StatusOK, gin.H{"success": false, "message": "Can't unlock"})
			return
		}
		p.IsLocked = false
		message = "Unlocked"
	} else {
		p.IsLocked = true
		p.PassphraseToUnlock = HashPassword(json.Passphrase)
		message = "Locked"
	}
	p.Save()
	c.JSON(http.StatusOK, gin.H{"success": true, "message": message})
}

func handleEncrypt(c *gin.Context) {
	type QueryJSON struct {
		Page       string `json:"page"`
		Passphrase string `json:"passphrase"`
	}

	var json QueryJSON
	if c.BindJSON(&json) != nil {
		c.String(http.StatusBadRequest, "Problem binding keys")
		return
	}
	p := Open(json.Page)
	if p.IsLocked {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "Locked"})
		return
	}
	q := Open(json.Page)
	var message string
	if p.IsEncrypted {
		decrypted, err2 := DecryptString(p.Text.GetCurrent(), json.Passphrase)
		if err2 != nil {
			c.JSON(http.StatusOK, gin.H{"success": false, "message": "Wrong password"})
			return
		}
		q.Erase()
		q = Open(json.Page)
		q.Update(decrypted)
		q.IsEncrypted = false
		q.IsLocked = p.IsLocked
		q.IsPrimedForSelfDestruct = p.IsPrimedForSelfDestruct
		message = "Decrypted"
	} else {
		currentText := p.Text.GetCurrent()
		encrypted, _ := EncryptString(currentText, json.Passphrase)
		q.Erase()
		q = Open(json.Page)
		q.Update(encrypted)
		q.IsEncrypted = true
		q.IsLocked = p.IsLocked
		q.IsPrimedForSelfDestruct = p.IsPrimedForSelfDestruct
		message = "Encrypted"
	}
	q.Save()
	c.JSON(http.StatusOK, gin.H{"success": true, "message": message})
}

func deleteListItem(c *gin.Context) {
	lineNum, err := strconv.Atoi(c.DefaultQuery("lineNum", "None"))
	page := c.Query("page") // shortcut for c.Request.URL.Query().Get("lastname")
	if err == nil {
		p := Open(page)

		_, listItems := reorderList(p.Text.GetCurrent())
		newText := p.Text.GetCurrent()
		for i, lineString := range listItems {
			// fmt.Println(i, lineString, lineNum)
			if i+1 == lineNum {
				// fmt.Println("MATCHED")
				if strings.Contains(lineString, "~~") == false {
					// fmt.Println(p.Text, "("+lineString[2:]+"\n"+")", "~~"+lineString[2:]+"~~"+"\n")
					newText = strings.Replace(newText+"\n", lineString[2:]+"\n", "~~"+strings.TrimSpace(lineString[2:])+"~~"+"\n", 1)
				} else {
					newText = strings.Replace(newText+"\n", lineString[2:]+"\n", lineString[4:len(lineString)-2]+"\n", 1)
				}
				p.Update(newText)
				break
			}
		}

		c.JSON(200, gin.H{
			"success": true,
			"message": "Done.",
		})
	} else {
		c.JSON(200, gin.H{
			"success": false,
			"message": err.Error(),
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/mattbaird/elastigo/api"
	"github.com/mattbaird/elastigo/core"
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"strings"
)

func handleHelp() string {
	return `
	Help
	----
	Commands:
	  eg:
	  host localhost
	  port 9200
	  index movies
	  get _search?q=title:thx1138
	`
}

func handleVersion() string {
	return `
	elRepl version 0.1
	`
}

func handleExit() string {
	fmt.Println("Bye.")
	os.Exit(0)
	return ""
}

func handleUnknownEntry(cmd *Command) string {
	return fmt.Sprintf("Command not found: %s", cmd.Name)
}

func handleHostSet(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandServer+" ")
	arg := cmd.Args
	server.host = arg
	return "Set server host: " + arg
}

func handleHostGet() string {
	return "Server host: " + server.host
}

func handlePortSet(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandPort+" ")
	arg := cmd.Args
	server.port = arg
	return "Set server port: " + arg
}

func handlePortGet() string {
	return "Server port: " + server.port
}

func handleIndexSet(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandIndex+" ")
	arg := cmd.Args
	server.index = arg
	return "Set index: " + arg
}

func handleIndexGet() string {
	return "Index: " + server.index
}

func handleDir(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandDir+" ")
	arg := cmd.Args
	if arg == "" {
		arg = "."
	}
	dirFiles, err := ioutil.ReadDir(arg)
	if err != nil {
		return err.Error()
	}
	files := ""
	for _, j := range dirFiles {
		files += j.Name() + "\n"
	}
	return files
}

func handleLog(cmd *Command) string {
	logLevel = 1
	return "Logging level set to: " + strconv.Itoa(logLevel)
}

func handleLoad(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandLoad+" ")
	arg := cmd.Args

	file, err := ioutil.ReadFile(arg)
	if err != nil {
		return err.Error()
	}
	fileText := string(file)
	loadedRequest.request = fileText
	return fileText
}

func handleRun(cmd *Command) string {
	//arg := cmd.Args
	loadedParts := strings.SplitN(loadedRequest.request, "\n", 2)

	loadedCmdParts := strings.SplitN(loadedParts[0], " ", 2)
	loadedCmd := loadedCmdParts[0]
	//loadedArgs := loadedCmdParts[1]

	//loadedQuery := loadedParts[1]

	cmdParser := NewCommandParser()
	newCmd, err := cmdParser.Parse(loadedRequest.request)
	if err != nil {
		return "Unable to parse loaded query for run command."
	}
	if strings.ToLower(loadedCmd) == "post" {
		resp := handlePost(newCmd)
		return resp
	} else if strings.ToLower(loadedCmd) == "put" {
		resp := handlePut(newCmd)
		return resp
	} else if strings.ToLower(loadedCmd) == "get" {
		resp := handleGet(newCmd)
		return resp
	}
	return "Unable to run loaded query."
}

func handleGet(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandGet+" ")
	arg := cmd.Args

	url := ""
	if server.index == "" {
		url = fmt.Sprintf("http://%s:%s/%s", server.host, server.port, arg)
	} else {
		url = fmt.Sprintf("http://%s:%s/%s/%s", server.host, server.port, server.index, arg)
	}

	fmt.Println("Request:", url)
	res, err := getHttpResource(url)
	if err != nil {
		return err.Error()
	}
	return res
}

// curl -XPUT "http://localhost:9200/movies/movie/1" -d'{ ... body ... }''
// becomes
// put movie/1 { "title": "Alien", "director": "Ridley Scott", "year": 1979, "genres": ["Science fiction"] }
// Currently, must be on single line.
func handlePut(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandPut+" ")
	arg := cmd.Args

	bodyIdx := strings.Index(arg, " ")
	queryArgs := arg[:bodyIdx]
	body := arg[bodyIdx:]

	url := fmt.Sprintf("http://%s:%s/%s/%s", server.host, server.port, server.index, queryArgs)
	fmt.Println("Request:", url)
	res, err := putHttpResource(url, body)
	if err != nil {
		return err.Error()
	}
	return res
}

// curl -XPOST "http://localhost:9200/movies/_search?pretty" -d'{ ... body ... }''
// becomes
// post _search?pretty { "query": { "term": { "director": "scott" } } }
// Currently, must be on single line.
func handlePost(cmd *Command) string {
	queryHost := server.host
	queryPort := server.port

	arg := cmd.Args
	bodyIdx := strings.Index(arg, "{")
	queryArgs := arg[:bodyIdx]
	queryArgs = strings.TrimPrefix(queryArgs, "/")
	queryArgs = strings.TrimSpace(queryArgs)
	body := arg[bodyIdx:]

	url := fmt.Sprintf("http://%s:%s/%s", queryHost, queryPort, queryArgs)

	fmt.Println("Request:", url)
	res, err := postHttpResource(url, body)
	if err != nil {
		return err.Error()
	}
	return res
}

// reindex localhost:9200/srcindex/type localhost:9200/targetindex/routing
func handleReindex(cmd *Command) string {
	fmt.Println("Reindexing...")
	//args := strings.TrimPrefix(entry, CommandReindex+" ")
	args := cmd.Args

	// \w+|"[\w\s]*"
	r, err := regexp.Compile(`^(.*?):(\d+?)/(.*?)/(.*?)/? (.*?):(\d+?)/(.*?)(/(.*?))?$`)
	if err != nil {
		return err.Error()
	}
	fmt.Println("Parsing command...")
	matches := r.FindAllStringSubmatch(args, -1)[0]
	fmt.Println("Parsed matches:", len(matches))
	srcHost := matches[1]
	srcPort := matches[2]
	srcIndex := matches[3]
	srcType := matches[4]
	tgtHost := matches[5]
	tgtPort := matches[6]
	tgtIndex := matches[7]
	tgtRouting := matches[8]

	api.Domain = srcHost
	api.Port = srcPort

	fmt.Println("Scanning...")
	scanArgs := map[string]interface{}{"search_type": "scan", "scroll": "1m", "size": "1000"}
	scanResult, err := core.SearchUri(srcIndex, srcType, scanArgs)
	if err != nil {
		fmt.Println("Failed getting scan result for index:", srcIndex, "; err:", err)
		return err.Error()
	}

	//total := scanResult.Hits.Total

	scrollId := scanResult.ScrollId
	counter := 0
	failures := 0

	fmt.Println("Scrolling...")
	scrollArgs := map[string]interface{}{"scroll": "1m"}
	scrollResult, err := core.Scroll(scrollArgs, scrollId)
	if err != nil {
		fmt.Println("Failed getting scroll result for index:", srcIndex, "; err:", err)
		return err.Error()
	}

	fmt.Println("Indexing...")
	var indexArgs map[string]interface{} = nil
	if tgtRouting != "" {
		indexArgs = map[string]interface{}{"routing": tgtRouting}
	}
	for len(scrollResult.Hits.Hits) > 0 {
		fmt.Println("Scroll result hits:", len(scrollResult.Hits.Hits))
		for _, j := range scrollResult.Hits.Hits {
			api.Domain = tgtHost
			api.Port = tgtPort
			_, err := core.Index(tgtIndex, srcType, j.Id, indexArgs, j.Source)
			if err != nil {
				fmt.Println("Failed inserting document, id:", j.Id, "; ", err)
				failures++
				continue
			}
			counter++
		}

		api.Domain = srcHost
		api.Port = srcPort
		// ScrollId changes with every request.
		scrollId = scrollResult.ScrollId
		scrollArgs := map[string]interface{}{"scroll": "1m"}
		scrollResult, err = core.Scroll(scrollArgs, scrollId)
		if err != nil {
			fmt.Println("Failed getting scroll result for index:", srcIndex, "; err:", err)
			return err.Error()
		}
	}
	return fmt.Sprintf("Total processed: %d.  %d failed.", counter, failures)
}
<commit_msg>Fixed body index not found in args.<commit_after>package main

import (
	"fmt"
	"github.com/mattbaird/elastigo/api"
	"github.com/mattbaird/elastigo/core"
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"strings"
)

func handleHelp() string {
	return `
	Help
	----
	Commands:
	  eg:
	  host localhost
	  port 9200
	  index movies
	  get _search?q=title:thx1138
	`
}

func handleVersion() string {
	return `
	elRepl version 0.1
	`
}

func handleExit() string {
	fmt.Println("Bye.")
	os.Exit(0)
	return ""
}

func handleUnknownEntry(cmd *Command) string {
	return fmt.Sprintf("Command not found: %s", cmd.Name)
}

func handleHostSet(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandServer+" ")
	arg := cmd.Args
	server.host = arg
	return "Set server host: " + arg
}

func handleHostGet() string {
	return "Server host: " + server.host
}

func handlePortSet(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandPort+" ")
	arg := cmd.Args
	server.port = arg
	return "Set server port: " + arg
}

func handlePortGet() string {
	return "Server port: " + server.port
}

func handleIndexSet(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandIndex+" ")
	arg := cmd.Args
	server.index = arg
	return "Set index: " + arg
}

func handleIndexGet() string {
	return "Index: " + server.index
}

func handleDir(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandDir+" ")
	arg := cmd.Args
	if arg == "" {
		arg = "."
	}
	dirFiles, err := ioutil.ReadDir(arg)
	if err != nil {
		return err.Error()
	}
	files := ""
	for _, j := range dirFiles {
		files += j.Name() + "\n"
	}
	return files
}

func handleLog(cmd *Command) string {
	logLevel = 1
	return "Logging level set to: " + strconv.Itoa(logLevel)
}

func handleLoad(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandLoad+" ")
	arg := cmd.Args

	file, err := ioutil.ReadFile(arg)
	if err != nil {
		return err.Error()
	}
	fileText := string(file)
	loadedRequest.request = fileText
	return fileText
}

func handleRun(cmd *Command) string {
	//arg := cmd.Args
	loadedParts := strings.SplitN(loadedRequest.request, "\n", 2)

	loadedCmdParts := strings.SplitN(loadedParts[0], " ", 2)
	loadedCmd := loadedCmdParts[0]
	//loadedArgs := loadedCmdParts[1]

	//loadedQuery := loadedParts[1]

	cmdParser := NewCommandParser()
	newCmd, err := cmdParser.Parse(loadedRequest.request)
	if err != nil {
		return "Unable to parse loaded query for run command."
	}
	if strings.ToLower(loadedCmd) == "post" {
		resp := handlePost(newCmd)
		return resp
	} else if strings.ToLower(loadedCmd) == "put" {
		resp := handlePut(newCmd)
		return resp
	} else if strings.ToLower(loadedCmd) == "get" {
		resp := handleGet(newCmd)
		return resp
	}
	return "Unable to run loaded query."
}

func handleGet(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandGet+" ")
	arg := cmd.Args

	url := ""
	if server.index == "" {
		url = fmt.Sprintf("http://%s:%s/%s", server.host, server.port, arg)
	} else {
		url = fmt.Sprintf("http://%s:%s/%s/%s", server.host, server.port, server.index, arg)
	}

	fmt.Println("Request:", url)
	res, err := getHttpResource(url)
	if err != nil {
		return err.Error()
	}
	return res
}

// curl -XPUT "http://localhost:9200/movies/movie/1" -d'{ ... body ... }''
// becomes
// put movie/1 { "title": "Alien", "director": "Ridley Scott", "year": 1979, "genres": ["Science fiction"] }
// Currently, must be on single line.
func handlePut(cmd *Command) string {
	//arg := strings.TrimPrefix(entry, CommandPut+" ")
	arg := cmd.Args

	bodyIdx := strings.Index(arg, " ")
	queryArgs := ""
	if bodyIdx > -1 {
		queryArgs = arg[:bodyIdx]
	}
	body := ""
	if bodyIdx > -1 {
		body = arg[bodyIdx:]
	}

	url := fmt.Sprintf("http://%s:%s/%s/%s", server.host, server.port, server.index, queryArgs)
	fmt.Println("Request:", url)
	res, err := putHttpResource(url, body)
	if err != nil {
		return err.Error()
	}
	return res
}

// curl -XPOST "http://localhost:9200/movies/_search?pretty" -d'{ ... body ... }''
// becomes
// post _search?pretty { "query": { "term": { "director": "scott" } } }
// Currently, must be on single line.
func handlePost(cmd *Command) string {
	queryHost := server.host
	queryPort := server.port

	arg := cmd.Args
	bodyIdx := strings.Index(arg, "{")
	queryArgs := arg[:bodyIdx]
	queryArgs = strings.TrimPrefix(queryArgs, "/")
	queryArgs = strings.TrimSpace(queryArgs)
	body := arg[bodyIdx:]

	url := fmt.Sprintf("http://%s:%s/%s", queryHost, queryPort, queryArgs)

	fmt.Println("Request:", url)
	res, err := postHttpResource(url, body)
	if err != nil {
		return err.Error()
	}
	return res
}

// reindex localhost:9200/srcindex/type localhost:9200/targetindex/routing
func handleReindex(cmd *Command) string {
	fmt.Println("Reindexing...")
	//args := strings.TrimPrefix(entry, CommandReindex+" ")
	args := cmd.Args

	// \w+|"[\w\s]*"
	r, err := regexp.Compile(`^(.*?):(\d+?)/(.*?)/(.*?)/? (.*?):(\d+?)/(.*?)(/(.*?))?$`)
	if err != nil {
		return err.Error()
	}
	fmt.Println("Parsing command...")
	matches := r.FindAllStringSubmatch(args, -1)[0]
	fmt.Println("Parsed matches:", len(matches))
	srcHost := matches[1]
	srcPort := matches[2]
	srcIndex := matches[3]
	srcType := matches[4]
	tgtHost := matches[5]
	tgtPort := matches[6]
	tgtIndex := matches[7]
	tgtRouting := matches[8]

	api.Domain = srcHost
	api.Port = srcPort

	fmt.Println("Scanning...")
	scanArgs := map[string]interface{}{"search_type": "scan", "scroll": "1m", "size": "1000"}
	scanResult, err := core.SearchUri(srcIndex, srcType, scanArgs)
	if err != nil {
		fmt.Println("Failed getting scan result for index:", srcIndex, "; err:", err)
		return err.Error()
	}

	//total := scanResult.Hits.Total

	scrollId := scanResult.ScrollId
	counter := 0
	failures := 0

	fmt.Println("Scrolling...")
	scrollArgs := map[string]interface{}{"scroll": "1m"}
	scrollResult, err := core.Scroll(scrollArgs, scrollId)
	if err != nil {
		fmt.Println("Failed getting scroll result for index:", srcIndex, "; err:", err)
		return err.Error()
	}

	fmt.Println("Indexing...")
	var indexArgs map[string]interface{} = nil
	if tgtRouting != "" {
		indexArgs = map[string]interface{}{"routing": tgtRouting}
	}
	for len(scrollResult.Hits.Hits) > 0 {
		fmt.Println("Scroll result hits:", len(scrollResult.Hits.Hits))
		for _, j := range scrollResult.Hits.Hits {
			api.Domain = tgtHost
			api.Port = tgtPort

			_, err := core.Index(tgtIndex, srcType, j.Id, indexArgs, j.Source)
			if err != nil {
				fmt.Println("Failed inserting document, id:", j.Id, "; ", err)
				failures++
				continue
			}
			counter++
		}

		api.Domain = srcHost
		api.Port = srcPort
		// ScrollId changes with every request.
		scrollId = scrollResult.ScrollId
		scrollArgs := map[string]interface{}{"scroll": "1m"}
		scrollResult, err = core.Scroll(scrollArgs, scrollId)
		if err != nil {
			fmt.Println("Failed getting scroll result for index:", srcIndex, "; err:", err)
			return err.Error()
		}
	}
	return fmt.Sprintf("Total processed: %d.  %d failed.", counter, failures)
}
<|endoftext|>


<commit_before>package gsproxy

import (
	"bytes"
	"sync"

	"github.com/gsdocker/gslogger"
	"github.com/gsrpc/gorpc"
	gorpcHandler "github.com/gsrpc/gorpc/handler"
)

type _TunnelServerHandler struct {
	gslogger.Log         // mixin log APIs
	proxy        *_Proxy // proxy
	id           byte    // agnet id
}

func (proxy *_Proxy) newTunnelServer() gorpc.Handler {
	return &_TunnelServerHandler{
		Log:   gslogger.Get("agent-server-tunnel"),
		proxy: proxy,
		id:    proxy.tunnelID(),
	}
}

func (handler *_TunnelServerHandler) Register(context gorpc.Context) error {
	return nil
}

func (handler *_TunnelServerHandler) Active(context gorpc.Context) error {
	return nil
}

func (handler *_TunnelServerHandler) Unregister(context gorpc.Context) {

}

func (handler *_TunnelServerHandler) Inactive(context gorpc.Context) {

}

func (handler *_TunnelServerHandler) CloseHandler(context gorpc.Context) {

}

func (handler *_TunnelServerHandler) MessageReceived(context gorpc.Context, message *gorpc.Message) (*gorpc.Message, error) {

	if message.Code != gorpc.CodeTunnel {
		return message, nil
	}

	handler.V("backward tunnel message")

	tunnel, err := gorpc.ReadTunnel(bytes.NewBuffer(message.Content))

	if err != nil {
		handler.E("backward tunnel(%s) message -- failed\n%s", tunnel.ID, err)
		return nil, err
	}

	if device, ok := handler.proxy.client(tunnel.ID); ok {

		tunnel.Message.Agent = handler.id

		err := device.SendMessage(tunnel.Message)

		if err == nil {
			handler.V("backward tunnel message -- success")
			return nil, nil
		}

		return nil, err
	}

	handler.E("backward tunnel(%s) message -- failed,device not found", tunnel.ID)

	return nil, nil
}

func (handler *_TunnelServerHandler) MessageSending(context gorpc.Context, message *gorpc.Message) (*gorpc.Message, error) {

	return message, nil
}

func (handler *_TunnelServerHandler) Panic(context gorpc.Context, err error) {

}

func (handler *_TunnelServerHandler) ID() byte {
	return handler.id
}

type _TransProxyHandler struct {
	gslogger.Log                   // mixin log APIs
	sync.RWMutex                   // mixin rw locker
	proxy        *_Proxy           // proxy
	client       *_Client          // client
	device       *gorpc.Device     // devices
	servers      map[uint16]Server // bound servers
	tunnels      map[byte]Server   // bound servers
}

func (proxy *_Proxy) newTransProxyHandler() gorpc.Handler {
	return &_TransProxyHandler{
		Log:     gslogger.Get("trans-proxy"),
		proxy:   proxy,
		servers: make(map[uint16]Server),
		tunnels: make(map[byte]Server),
	}
}

func (handler *_TransProxyHandler) bind(id uint16, server Server) {
	handler.Lock()
	defer handler.Unlock()

	tunnel, _ := server.Handler(tunnelHandler)

	handler.servers[id] = server

	handler.tunnels[tunnel.(*_TunnelServerHandler).ID()] = server
}

func (handler *_TransProxyHandler) unbind(id uint16) {
	handler.Lock()
	defer handler.Unlock()

	delete(handler.servers, id)
}

func (handler *_TransProxyHandler) Register(context gorpc.Context) error {
	return nil
}

func (handler *_TransProxyHandler) Active(context gorpc.Context) error {

	dh, _ := context.Pipeline().Handler(dhHandler)

	handler.device = dh.(gorpcHandler.CryptoServer).GetDevice()

	return nil
}

func (handler *_TransProxyHandler) Unregister(context gorpc.Context) {

}

func (handler *_TransProxyHandler) Inactive(context gorpc.Context) {

}

func (handler *_TransProxyHandler) forward(server Server, message *gorpc.Message) error {
	handler.V("forward tunnel(%s) message", handler.device)

	tunnel := gorpc.NewTunnel()

	tunnel.ID = handler.device

	tunnel.Message = message

	var buff bytes.Buffer

	err := gorpc.WriteTunnel(&buff, tunnel)

	if err != nil {
		return err
	}

	message.Code = gorpc.CodeTunnel

	message.Content = buff.Bytes()

	err = server.SendMessage(message)

	if err == nil {
		handler.V("forward tunnel(%s) message(%p) -- success", handler.device, message)
	} else {
		handler.E("forward tunnel(%s) message -- failed\n%s", handler.device, err)
	}

	return err
}

func (handler *_TransProxyHandler) MessageReceived(context gorpc.Context, message *gorpc.Message) (*gorpc.Message, error) {

	if message.Code == gorpc.CodeResponse {

		if server, ok := handler.tunnels[message.Agent]; ok {

			return nil, handler.forward(server, message)
		}

		return message, nil

	}

	if message.Code != gorpc.CodeRequest {
		return message, nil
	}

	request, err := gorpc.ReadRequest(bytes.NewBuffer(message.Content))

	if err != nil {
		handler.E("[%s] unmarshal request error\n%s", handler.proxy.name, err)
		return nil, err
	}

	service := request.Service

	if server, ok := handler.servers[service]; ok {

		handler.V("forward tunnel(%s) message", handler.device)

		tunnel := gorpc.NewTunnel()

		tunnel.ID = handler.device

		tunnel.Message = message

		var buff bytes.Buffer

		err := gorpc.WriteTunnel(&buff, tunnel)

		if err != nil {
			return nil, err
		}

		message.Code = gorpc.CodeTunnel

		message.Content = buff.Bytes()

		err = server.SendMessage(message)

		if err == nil {
			handler.V("forward tunnel(%s) message(%p) -- success", handler.device, message)
		} else {
			handler.E("forward tunnel(%s) message -- failed\n%s", handler.device, err)
		}

		return nil, err
	}

	return message, nil
}

func (handler *_TransProxyHandler) MessageSending(context gorpc.Context, message *gorpc.Message) (*gorpc.Message, error) {

	return message, nil
}

func (handler *_TransProxyHandler) Panic(context gorpc.Context, err error) {

}
<commit_msg>fix reconnect bug<commit_after>package gsproxy

import (
	"bytes"
	"sync"

	"github.com/gsdocker/gslogger"
	"github.com/gsrpc/gorpc"
	gorpcHandler "github.com/gsrpc/gorpc/handler"
)

type _TunnelServerHandler struct {
	gslogger.Log         // mixin log APIs
	proxy        *_Proxy // proxy
	id           byte    // agnet id
}

func (proxy *_Proxy) newTunnelServer() gorpc.Handler {
	return &_TunnelServerHandler{
		Log:   gslogger.Get("agent-server-tunnel"),
		proxy: proxy,
		id:    proxy.tunnelID(),
	}
}

func (handler *_TunnelServerHandler) Register(context gorpc.Context) error {
	return nil
}

func (handler *_TunnelServerHandler) Active(context gorpc.Context) error {
	return nil
}

func (handler *_TunnelServerHandler) Unregister(context gorpc.Context) {

}

func (handler *_TunnelServerHandler) Inactive(context gorpc.Context) {

}

func (handler *_TunnelServerHandler) CloseHandler(context gorpc.Context) {

}

func (handler *_TunnelServerHandler) MessageReceived(context gorpc.Context, message *gorpc.Message) (*gorpc.Message, error) {

	if message.Code != gorpc.CodeTunnel {
		return message, nil
	}

	handler.V("backward tunnel message")

	tunnel, err := gorpc.ReadTunnel(bytes.NewBuffer(message.Content))

	if err != nil {
		handler.E("backward tunnel(%s) message -- failed\n%s", tunnel.ID, err)
		return nil, err
	}

	if device, ok := handler.proxy.client(tunnel.ID); ok {

		tunnel.Message.Agent = handler.id

		err := device.SendMessage(tunnel.Message)

		if err == nil {
			handler.V("backward tunnel message -- success")
			return nil, nil
		}

		return nil, err
	}

	handler.E("backward tunnel(%s) message -- failed,device not found", tunnel.ID)

	return nil, nil
}

func (handler *_TunnelServerHandler) MessageSending(context gorpc.Context, message *gorpc.Message) (*gorpc.Message, error) {

	return message, nil
}

func (handler *_TunnelServerHandler) Panic(context gorpc.Context, err error) {

}

func (handler *_TunnelServerHandler) ID() byte {
	return handler.id
}

type _TransProxyHandler struct {
	gslogger.Log                   // mixin log APIs
	sync.RWMutex                   // mixin rw locker
	proxy        *_Proxy           // proxy
	client       *_Client          // client
	device       *gorpc.Device     // devices
	servers      map[uint16]Server // bound servers
	tunnels      map[byte]Server   // bound servers
}

func (proxy *_Proxy) newTransProxyHandler() gorpc.Handler {
	return &_TransProxyHandler{
		Log:     gslogger.Get("trans-proxy"),
		proxy:   proxy,
		servers: make(map[uint16]Server),
		tunnels: make(map[byte]Server),
	}
}

func (handler *_TransProxyHandler) bind(id uint16, server Server) {
	handler.Lock()
	defer handler.Unlock()

	tunnel, _ := server.Handler(tunnelHandler)

	handler.servers[id] = server

	handler.tunnels[tunnel.(*_TunnelServerHandler).ID()] = server
}

func (handler *_TransProxyHandler) unbind(id uint16) {
	handler.Lock()
	defer handler.Unlock()

	delete(handler.servers, id)
}

func (handler *_TransProxyHandler) Register(context gorpc.Context) error {
	return nil
}

func (handler *_TransProxyHandler) Active(context gorpc.Context) error {

	dh, _ := context.Pipeline().Handler(dhHandler)

	handler.device = dh.(gorpcHandler.CryptoServer).GetDevice()

	return nil
}

func (handler *_TransProxyHandler) Unregister(context gorpc.Context) {

}

func (handler *_TransProxyHandler) Inactive(context gorpc.Context) {

}

func (handler *_TransProxyHandler) forward(server Server, message *gorpc.Message) error {
	handler.V("forward tunnel(%s) message", handler.device)

	tunnel := gorpc.NewTunnel()

	tunnel.ID = handler.device

	tunnel.Message = message

	var buff bytes.Buffer

	err := gorpc.WriteTunnel(&buff, tunnel)

	if err != nil {
		return err
	}

	message.Code = gorpc.CodeTunnel

	message.Content = buff.Bytes()

	err = server.SendMessage(message)

	if err == nil {
		handler.V("forward tunnel(%s) message(%p) -- success", handler.device, message)
	} else {
		handler.E("forward tunnel(%s) message -- failed\n%s", handler.device, err)
	}

	return err
}

func (handler *_TransProxyHandler) tunnel(agent byte) (Server, bool) {

	handler.RLock()
	defer handler.RUnlock()

	server, ok := handler.tunnels[agent]

	return server, ok
}

func (handler *_TransProxyHandler) server(service uint16) (Server, bool) {

	handler.RLock()
	defer handler.RUnlock()

	server, ok := handler.servers[service]

	return server, ok
}

func (handler *_TransProxyHandler) MessageReceived(context gorpc.Context, message *gorpc.Message) (*gorpc.Message, error) {

	if message.Code == gorpc.CodeResponse {

		if server, ok := handler.tunnel(message.Agent); ok {

			err := handler.forward(server, message)

			if err != nil {
				context.Close()
			}

			return nil, err
		}

		return message, nil

	}

	if message.Code != gorpc.CodeRequest {
		return message, nil
	}

	request, err := gorpc.ReadRequest(bytes.NewBuffer(message.Content))

	if err != nil {
		handler.E("[%s] unmarshal request error\n%s", handler.proxy.name, err)
		return nil, err
	}

	service := request.Service

	if server, ok := handler.server(service); ok {

		handler.V("forward tunnel(%s) message", handler.device)

		tunnel := gorpc.NewTunnel()

		tunnel.ID = handler.device

		tunnel.Message = message

		var buff bytes.Buffer

		err := gorpc.WriteTunnel(&buff, tunnel)

		if err != nil {
			return nil, err
		}

		message.Code = gorpc.CodeTunnel

		message.Content = buff.Bytes()

		err = server.SendMessage(message)

		if err != nil {
			context.Close()
			handler.V("forward tunnel(%s) message(%p) -- failed\n%s", handler.device, message, err)
			return nil, err
		}

		handler.V("forward tunnel(%s) message(%p) -- success", handler.device, message)

		return nil, err
	}

	return message, nil
}

func (handler *_TransProxyHandler) MessageSending(context gorpc.Context, message *gorpc.Message) (*gorpc.Message, error) {

	return message, nil
}

func (handler *_TransProxyHandler) Panic(context gorpc.Context, err error) {

}
<|endoftext|>