<commit_before>package rerun

import (
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/fsnotify/fsnotify"
)

type Watcher struct {
	watcher *fsnotify.Watcher
	watch   map[string]struct{}
	ignore  map[string]struct{}
	done    chan struct{}
}

type ChangeSet struct {
	Files map[string]struct{}
	Error error
}

func NewWatcher() (*Watcher, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}

	w := &Watcher{
		watcher: watcher,
		watch:   make(map[string]struct{}),
		ignore:  make(map[string]struct{}),
		done:    make(chan struct{}, 0),
	}

	return w, nil
}

func (w *Watcher) Add(paths ...string) {
	for _, path := range paths {
		w.watch[path] = struct{}{}
		//fmt.Printf("Add %v\n", path)
	}
}

func (w *Watcher) Ignore(paths ...string) {
	for _, path := range paths {
		w.ignore[path] = struct{}{}
		//fmt.Printf("Ignore %v\n", path)
	}
}

func (w *Watcher) Watch(delay time.Duration) <-chan ChangeSet {
	//	fmt.Println()

	// resolve add + ignore paths
	for path, _ := range w.watch { //s
		filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
			if !info.IsDir() {
				return nil
			}
			if strings.HasSuffix(path, ".git") {
				//fmt.Printf("skip %v\n", path)
				return filepath.SkipDir
			}

			if _, ok := w.ignore[path]; ok {
				//fmt.Printf("skip %v\n", path)
				return filepath.SkipDir
			}

			w.watcher.Add(path)
			//fmt.Printf("watch %v\n", path)
			return nil
		})
	}
	//	fmt.Println()

	changes := make(chan ChangeSet, 1)

	go func() {
		for {
			change := ChangeSet{
				Files: make(map[string]struct{}),
			}

			timeout := time.NewTimer(1<<63 - 1) // max duration
			timeout.Stop()

		loop:
			for {
				select {
				case event := <-w.watcher.Events:
					// Ignore CHMOD.
					if event.Op&fsnotify.Chmod == fsnotify.Chmod {
						continue
					}

					timeout.Reset(delay)

					//fmt.Printf("event: %v (%v)\n", event, time.Now()) //
					// if event.Op&fsnotify.Write == fsnotify.Write {
					// 	log.Println("modified file:", event.Name)
					// }
					change.Files[event.Name] = struct{}{}

				case err := <-w.watcher.Errors:
					change.Error = err
					changes <- change
					timeout.Stop()
					break loop

				case <-timeout.C:
					changes <- change
					timeout.Stop()
					break loop

				case <-w.done:
					close(changes)
					timeout.Stop()
					return
				}
			}
		}
	}()

	return changes
}

func (w *Watcher) Close() error {
	close(w.done)
	return w.watcher.Close()
}

func (c *ChangeSet) String() string {
	str := ""
	for file, _ := range c.Files {
		str += "\n" + file
	}
	return str
}
<commit_msg>Fix panic on empty nil pointer dereference<commit_after>package rerun

import (
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/fsnotify/fsnotify"
)

type Watcher struct {
	watcher *fsnotify.Watcher
	watch   map[string]struct{}
	ignore  map[string]struct{}
	done    chan struct{}
}

type ChangeSet struct {
	Files map[string]struct{}
	Error error
}

func NewWatcher() (*Watcher, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}

	w := &Watcher{
		watcher: watcher,
		watch:   make(map[string]struct{}),
		ignore:  make(map[string]struct{}),
		done:    make(chan struct{}, 0),
	}

	return w, nil
}

func (w *Watcher) Add(paths ...string) {
	for _, path := range paths {
		w.watch[path] = struct{}{}
		//fmt.Printf("Add %v\n", path)
	}
}

func (w *Watcher) Ignore(paths ...string) {
	for _, path := range paths {
		w.ignore[path] = struct{}{}
		//fmt.Printf("Ignore %v\n", path)
	}
}

func (w *Watcher) Watch(delay time.Duration) <-chan ChangeSet {
	//	fmt.Println()

	// resolve add + ignore paths
	for path, _ := range w.watch { //s
		filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
			if info != nil && !info.IsDir() {
				return nil
			}
			if strings.HasSuffix(path, ".git") {
				//fmt.Printf("skip %v\n", path)
				return filepath.SkipDir
			}

			if _, ok := w.ignore[path]; ok {
				//fmt.Printf("skip %v\n", path)
				return filepath.SkipDir
			}

			w.watcher.Add(path)
			//fmt.Printf("watch %v\n", path)
			return nil
		})
	}
	//	fmt.Println()

	changes := make(chan ChangeSet, 1)

	go func() {
		for {
			change := ChangeSet{
				Files: make(map[string]struct{}),
			}

			timeout := time.NewTimer(1<<63 - 1) // max duration
			timeout.Stop()

		loop:
			for {
				select {
				case event := <-w.watcher.Events:
					// Ignore CHMOD.
					if event.Op&fsnotify.Chmod == fsnotify.Chmod {
						continue
					}

					timeout.Reset(delay)

					//fmt.Printf("event: %v (%v)\n", event, time.Now()) //
					// if event.Op&fsnotify.Write == fsnotify.Write {
					// 	log.Println("modified file:", event.Name)
					// }
					change.Files[event.Name] = struct{}{}

				case err := <-w.watcher.Errors:
					change.Error = err
					changes <- change
					timeout.Stop()
					break loop

				case <-timeout.C:
					changes <- change
					timeout.Stop()
					break loop

				case <-w.done:
					close(changes)
					timeout.Stop()
					return
				}
			}
		}
	}()

	return changes
}

func (w *Watcher) Close() error {
	close(w.done)
	return w.watcher.Close()
}

func (c *ChangeSet) String() string {
	str := ""
	for file, _ := range c.Files {
		str += "\n" + file
	}
	return str
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"github.com/scottdware/go-requestor"
	"os"
	"strings"
)

type CVE struct {
	Title       string `xml:"title"`
	Link        string `xml:"link"`
	Description string `xml:"description"`
	Date        string `xml:"date"`
}

type Vulnerabilities struct {
	XMLName xml.Name `xml:"RDF"`
	CVEs    []CVE    `xml:"item"`
}

var (
	rssNormal   = "https://nvd.nist.gov/download/nvd-rss.xml"
	rssAnalyzed = "https://nvd.nist.gov/download/nvd-rss-analyzed.xml"
	vulns       Vulnerabilities
	list        bool
	cve         string
	analyzed    bool
)

func init() {
	flag.Usage = func() {
		fmt.Println("cveget - List recent CVE's and view information about them.\n")
		fmt.Println("Usage: cveget [OPTIONS]\n")
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.BoolVar(&list, "list", false, "List all CVE's within the previous eight days.")
	flag.BoolVar(&analyzed, "analyzed", false, "Provides only vulnerabilities which have been analyzed within the previous eight days.")
	flag.StringVar(&cve, "cve", "", "Specify a CVE to view information on.")
	flag.Parse()

	if !list && cve == "" {
		flag.Usage()
	}
}

func main() {
	feedData := requestor.Send(rssNormal, nil)
	if feedData.Error != nil {
		fmt.Println(feedData.Error)
	}

	if analyzed {
		feedData = requestor.Send(rssAnalyzed, nil)
		if feedData.Error != nil {
			fmt.Println(feedData.Error)
		}
	}

	if err := xml.Unmarshal(feedData.Body, &vulns); err != nil {
		fmt.Println(err)
	}

	matches := len(vulns.CVEs)

	fmt.Printf("%d total CVE's\n\n", matches)

	for n, c := range vulns.CVEs {
		if list {
			fmt.Println(c.Title)
		}

		if !list {
			if strings.Contains(c.Title, cve) {
				fmt.Printf("%d) %s\n\n", n, c.Title)
				fmt.Printf("Published: %s\n", c.Date)
				fmt.Printf("%s\n\n", c.Link)
				fmt.Printf("%s\n\n", c.Description)
			}
		}
	}
}
<commit_msg>Removed CVE count<commit_after>package main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"github.com/scottdware/go-requestor"
	"os"
	"strings"
)

type CVE struct {
	Title       string `xml:"title"`
	Link        string `xml:"link"`
	Description string `xml:"description"`
	Date        string `xml:"date"`
}

type Vulnerabilities struct {
	XMLName xml.Name `xml:"RDF"`
	CVEs    []CVE    `xml:"item"`
}

var (
	rssNormal   = "https://nvd.nist.gov/download/nvd-rss.xml"
	rssAnalyzed = "https://nvd.nist.gov/download/nvd-rss-analyzed.xml"
	vulns       Vulnerabilities
	list        bool
	cve         string
	analyzed    bool
)

func init() {
	flag.Usage = func() {
		fmt.Println("cveget - List recent CVE's and view information about them.\n")
		fmt.Println("Usage: cveget [OPTIONS]\n")
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.BoolVar(&list, "list", false, "List all CVE's within the previous eight days.")
	flag.BoolVar(&analyzed, "analyzed", false, "Provides only vulnerabilities which have been analyzed within the previous eight days.")
	flag.StringVar(&cve, "cve", "", "Specify a CVE to view information on.")
	flag.Parse()

	if !list && cve == "" {
		flag.Usage()
	}
}

func main() {
	feedData := requestor.Send(rssNormal, nil)
	if feedData.Error != nil {
		fmt.Println(feedData.Error)
	}

	if analyzed {
		feedData = requestor.Send(rssAnalyzed, nil)
		if feedData.Error != nil {
			fmt.Println(feedData.Error)
		}
	}

	if err := xml.Unmarshal(feedData.Body, &vulns); err != nil {
		fmt.Println(err)
	}

	for _, c := range vulns.CVEs {
		if list {
			fmt.Println(c.Title)
		}

		if !list {
			if strings.Contains(c.Title, cve) {
				fmt.Printf("%s\n\n", c.Title)
				fmt.Printf("Published: %s\n", c.Date)
				fmt.Printf("%s\n\n", c.Link)
				fmt.Printf("%s\n\n", c.Description)
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Igor Dolzhikov. All rights reserved.
// Use of this source code is governed by
// license that can be found in the LICENSE file.

/*
Package daemon 0.2.2 for use with Go (golang) services.

Package daemon provides primitives for daemonization of golang services.
This package is not provide implementation of user daemon,
accordingly must have root rights to install/remove service.
In the current implementation is only supported Linux and Mac Os X daemon.

Example:

	//Example of the daemon with echo service
	package main

	import (
		"fmt"
		"log"
		"net"
		"os"
		"os/signal"
		"syscall"

		"github.com/takama/daemon"
	)

	const (

		// name of the service, match with executable file name
		name        = "myservice"
		description = "My Echo Service"

		// port which daemon should be listen
		port = ":9977"
	)

	// Service has embedded daemon
	type Service struct {
		daemon.Daemon
	}

	// Manage by daemon commands or run the daemon
	func (service *Service) Manage() (string, error) {

		usage := "Usage: myservice install | remove | start | stop | status"

		// if received any kind of command, do it
		if len(os.Args) > 1 {
			command := os.Args[1]
			switch command {
			case "install":
				return service.Install()
			case "remove":
				return service.Remove()
			case "start":
				return service.Start()
			case "stop":
				return service.Stop()
			case "status":
				return service.Status()
			default:
				return usage, nil
			}
		}

		// Do something, call your goroutines, etc

		// Set up channel on which to send signal notifications.
		// We must use a buffered channel or risk missing the signal
		// if we're not ready to receive when the signal is sent.
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

		// Set up listener for defined host and port
		listener, err := net.Listen("tcp", port)
		if err != nil {
			return "Possibly was a problem with the port binding", err
		}

		// set up channel on which to send accepted connections
		listen := make(chan net.Conn, 100)
		go acceptConnection(listener, listen)

		// loop work cycle with accept connections or interrupt
		// by system signal
		for {
			select {
			case conn := <-listen:
				go handleClient(conn)
			case killSignal := <-interrupt:
				log.Println("Got signal:", killSignal)
				log.Println("Stoping listening on ", listener.Addr())
				listener.Close()
				if killSignal == os.Interrupt {
					return "Daemon was interruped by system signal", nil
				}
				return "Daemon was killed", nil
			}
		}

		// never happen, but need to complete code
		return usage, nil
	}

	// Accept a client connection and collect it in a channel
	func acceptConnection(listener net.Listener, listen chan<- net.Conn) {
		for {
			conn, err := listener.Accept()
			if err != nil {
				continue
			}
			listen <- conn
		}
	}

	func handleClient(client net.Conn) {
		for {
			buf := make([]byte, 4096)
			numbytes, err := client.Read(buf)
			if numbytes == 0 || err != nil {
				return
			}
			client.Write(buf)
		}
	}

	func main() {
		srv, err := daemon.New(name, description)
		if err != nil {
			fmt.Println("Error: ", err)
			os.Exit(1)
		}
		service := &Service{srv}
		status, err := service.Manage()
		if err != nil {
			fmt.Println(status, "\nError: ", err)
			os.Exit(1)
		}
		fmt.Println(status)
	}

Go daemon
*/
package daemon

// Daemon interface has standard set of a methods/commands
type Daemon interface {

	// Install the service into the system
	Install() (string, error)

	// Remove the service and all corresponded files from the system
	Remove() (string, error)

	// Start the service
	Start() (string, error)

	// Stop the service
	Stop() (string, error)

	// Status - check the service status
	Status() (string, error)
}

// New - Create a new daemon
//
// name: name of the service, match with executable file name;
// description: any explanation, what is the service, its purpose
func New(name, description string) (Daemon, error) {
	return newDaemon(name, description)
}
<commit_msg>Bumped version number to 0.2.3<commit_after>// Copyright 2014 Igor Dolzhikov. All rights reserved.
// Use of this source code is governed by
// license that can be found in the LICENSE file.

/*
Package daemon 0.2.3 for use with Go (golang) services.

Package daemon provides primitives for daemonization of golang services.
This package is not provide implementation of user daemon,
accordingly must have root rights to install/remove service.
In the current implementation is only supported Linux and Mac Os X daemon.

Example:

	//Example of the daemon with echo service
	package main

	import (
		"fmt"
		"log"
		"net"
		"os"
		"os/signal"
		"syscall"

		"github.com/takama/daemon"
	)

	const (

		// name of the service, match with executable file name
		name        = "myservice"
		description = "My Echo Service"

		// port which daemon should be listen
		port = ":9977"
	)

	// Service has embedded daemon
	type Service struct {
		daemon.Daemon
	}

	// Manage by daemon commands or run the daemon
	func (service *Service) Manage() (string, error) {

		usage := "Usage: myservice install | remove | start | stop | status"

		// if received any kind of command, do it
		if len(os.Args) > 1 {
			command := os.Args[1]
			switch command {
			case "install":
				return service.Install()
			case "remove":
				return service.Remove()
			case "start":
				return service.Start()
			case "stop":
				return service.Stop()
			case "status":
				return service.Status()
			default:
				return usage, nil
			}
		}

		// Do something, call your goroutines, etc

		// Set up channel on which to send signal notifications.
		// We must use a buffered channel or risk missing the signal
		// if we're not ready to receive when the signal is sent.
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

		// Set up listener for defined host and port
		listener, err := net.Listen("tcp", port)
		if err != nil {
			return "Possibly was a problem with the port binding", err
		}

		// set up channel on which to send accepted connections
		listen := make(chan net.Conn, 100)
		go acceptConnection(listener, listen)

		// loop work cycle with accept connections or interrupt
		// by system signal
		for {
			select {
			case conn := <-listen:
				go handleClient(conn)
			case killSignal := <-interrupt:
				log.Println("Got signal:", killSignal)
				log.Println("Stoping listening on ", listener.Addr())
				listener.Close()
				if killSignal == os.Interrupt {
					return "Daemon was interruped by system signal", nil
				}
				return "Daemon was killed", nil
			}
		}

		// never happen, but need to complete code
		return usage, nil
	}

	// Accept a client connection and collect it in a channel
	func acceptConnection(listener net.Listener, listen chan<- net.Conn) {
		for {
			conn, err := listener.Accept()
			if err != nil {
				continue
			}
			listen <- conn
		}
	}

	func handleClient(client net.Conn) {
		for {
			buf := make([]byte, 4096)
			numbytes, err := client.Read(buf)
			if numbytes == 0 || err != nil {
				return
			}
			client.Write(buf)
		}
	}

	func main() {
		srv, err := daemon.New(name, description)
		if err != nil {
			fmt.Println("Error: ", err)
			os.Exit(1)
		}
		service := &Service{srv}
		status, err := service.Manage()
		if err != nil {
			fmt.Println(status, "\nError: ", err)
			os.Exit(1)
		}
		fmt.Println(status)
	}

Go daemon
*/
package daemon

// Daemon interface has standard set of a methods/commands
type Daemon interface {

	// Install the service into the system
	Install() (string, error)

	// Remove the service and all corresponded files from the system
	Remove() (string, error)

	// Start the service
	Start() (string, error)

	// Stop the service
	Stop() (string, error)

	// Status - check the service status
	Status() (string, error)
}

// New - Create a new daemon
//
// name: name of the service, match with executable file name;
// description: any explanation, what is the service, its purpose
func New(name, description string) (Daemon, error) {
	return newDaemon(name, description)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

/* This function takes a string and returns
   a (potentially nil) error object */
func TTWS(filename string) error {
	/* Open the input file */
	inf, err := os.Open(filename);
	/* In case this function generates a "panic", be sure to close this file */
	defer inf.Close();
	/* Did we open it successfully?  If not, close and return. */
	if (err!=nil) { inf.Close(); return err; }

	/* Open the output file */
	outf, err := os.Create(filename+"1");
	/* In case this function generates a "panic", be sure to close this file */
	defer outf.Close();
	/* Did we open it succesfully?  If not, close all and return. */
	if (err!=nil) { inf.Close(); outf.Close(); return err; }

	/* Create a scanner object to break this in to lines */
	scanner := bufio.NewScanner(inf)
	/* Declare a variable for the line */
	var line string;
	/* Loop over lines */
	for scanner.Scan() {
		/* Trim right space and then add the \n back on the end before writing */
		line = strings.TrimRight(scanner.Text(), " \t")+"\n"
		outf.Write([]byte(line));
	}

	/* Close all open files */
	inf.Close();
	outf.Close();

	/* No errors, so we return nil */
	return nil;
}

func WalkFunc(path string, fi os.FileInfo, err error) error {
	/* list of directories to ignore */
	blacklist := []string{".bzr", ".cvs", ".git", ".hg", ".svn"}
	if contains(path, blacklist){
		fmt.Printf("Skipping version control dir: %s\n", path)
		return filepath.SkipDir
	} else {
		inf, err := os.Open(path)
		defer inf.Close();
		if (err!=nil) { inf.Close(); return err; }
		readStart := io.LimitReader(inf, 512);
		data, err := ioutil.ReadAll(readStart);
		/* Close all open files */
		inf.Close();
		/* Determine file type */
		fileType := http.DetectContentType(data);
		if (strings.Contains(fileType, "text/plain") && !fi.IsDir()){
			fmt.Printf("Trimming: %v\n", path);
			TTWS(path);
		} else if !fi.IsDir() {
			fmt.Printf("Skipping file of type '%v': %v\n", fileType, path)
		}
	}
	return nil
}

func contains(x string, a []string) bool {
	for _, e := range(a) {
		if (x==e) { return true; }
	}
	return false;
}

func main() {
	flag.Parse()
	root := flag.Arg(0)
	err := filepath.Walk(root, WalkFunc)
	fmt.Printf("filepath.Walk() returned %v\n", err)
}
<commit_msg>Write out trimmed content to temp file and then replace the original.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

/* This function takes a string and returns
   a (potentially nil) error object */
func TTWS(filename string) error {
	/* Open the input file */
	inf, err := os.Open(filename);
	/* In case this function generates a "panic", be sure to close this file */
	defer inf.Close();
	/* Did we open it successfully?  If not, close and return. */
	if (err!=nil) { inf.Close(); return err; }

	/* Open the output file in system temp dir*/
	outf, err := ioutil.TempFile("","");
	/* In case this function generates a "panic", be sure to close this file */
	defer outf.Close();
	/* Did we open it succesfully?  If not, close all and return. */
	if (err!=nil) { inf.Close(); outf.Close(); return err; }
	/* Create a scanner object to break this in to lines */
	scanner := bufio.NewScanner(inf);
	/* Declare a variable for the line */
	var line string;
	/* Loop over lines */
	for scanner.Scan() {
		/* Trim right space and then add the \n back on the end before writing */
		line = strings.TrimRight(scanner.Text(), " \t")+"\n"
		outf.Write([]byte(line));
	}
	/* Close all open files */
	inf.Close();
	outf.Close();
	/* Replace the source file by the trimmed file */
	os.Rename(outf.Name(), filename);

	/* No errors, so we return nil */
	return nil;
}

func WalkFunc(path string, fi os.FileInfo, err error) error {
	/* list of directories to ignore */
	blacklist := []string{".bzr", ".cvs", ".git", ".hg", ".svn"}
	if contains(path, blacklist){
		fmt.Printf("Skipping version control dir: %s\n", path)
		return filepath.SkipDir
	} else {
		inf, err := os.Open(path)
		defer inf.Close();
		if (err!=nil) { inf.Close(); return err; }
		readStart := io.LimitReader(inf, 512);
		data, err := ioutil.ReadAll(readStart);
		/* Close all open files */
		inf.Close();
		/* Determine file type */
		fileType := http.DetectContentType(data);
		/* only act on text files */
		if (strings.Contains(fileType, "text/plain") && !fi.IsDir()){
			fmt.Printf("Trimming: %v\n", path);
			TTWS(path);
		} else if !fi.IsDir() {
			fmt.Printf("Skipping file of type '%v': %v\n", fileType, path)
		}
	}
	return nil
}

func contains(x string, a []string) bool {
	for _, e := range(a) {
		if (x==e) { return true; }
	}
	return false;
}

func main() {
	flag.Parse()
	root := flag.Arg(0)
	err := filepath.Walk(root, WalkFunc)
	fmt.Printf("filepath.Walk() returned %v\n", err)
}
<|endoftext|>


<commit_before>package handlers

import (
	"github.com/ant0ine/go-json-rest/rest"
	"go-message-masking/persistence"
	"net/http"
	"regexp"
)

type Message struct {
	Locale string
	Text   string
}

func MaskSensitiveData(w rest.ResponseWriter, r *rest.Request) {

	message := Message{}
	err := r.DecodeJsonPayload(&message)

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

	processed_message := mask_sensitive_data(message.Text, persistence.Expressions)
	w.WriteJson(
		&Message{
			Locale: message.Locale,
			Text:   processed_message,
		},
	)
}

func mask_sensitive_data(s string, expressions_map map[string]string) string {
	for _, value := range expressions_map {
		s = apply_expression(s, value)
	}

	return s
}

func apply_expression(s string, expression string) string {
	re := regexp.MustCompile(expression)
	return re.ReplaceAllStringFunc(s, func(str string) string {
		var mask string
		for range str {
			mask += "X"
		}
		return mask
	})
}
<commit_msg>refactor: Rename functions based on style guide recommendations<commit_after>package handlers

import (
	"go-message-masking/persistence"
	"net/http"
	"regexp"

	"github.com/ant0ine/go-json-rest/rest"
)

type Message struct {
	Locale string
	Text   string
}

func MaskSensitiveData(w rest.ResponseWriter, r *rest.Request) {

	message := Message{}
	err := r.DecodeJsonPayload(&message)

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

	processed_message := maskSensitiveData(message.Text, persistence.Expressions)
	w.WriteJson(
		&Message{
			Locale: message.Locale,
			Text:   processed_message,
		},
	)
}

func maskSensitiveData(s string, expressionMap map[string]string) string {
	for _, value := range expressionMap {
		s = applyExpression(s, value)
	}

	return s
}

func applyExpression(s string, expression string) string {
	re := regexp.MustCompile(expression)
	return re.ReplaceAllStringFunc(s, func(str string) string {
		var mask string
		for range str {
			mask += "X"
		}
		return mask
	})
}
<|endoftext|>


<commit_before>package stream

import (
	"context"
	"testing"
	time "time"

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

func noopUnSub() {}

func TestSubscription(t *testing.T) {
	eb := newEventBuffer()

	index := uint64(100)

	startHead := eb.Head()

	// Start with an event in the buffer
	publishTestEvent(index, eb, "test")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	req := SubscribeRequest{
		Topic: testTopic,
		Key:   "test",
	}
	sub := newSubscription(req, startHead, noopUnSub)

	// First call to sub.Next should return our published event immediately
	start := time.Now()
	got, err := sub.Next(ctx)
	elapsed := time.Since(start)
	require.NoError(t, err)
	require.True(t, elapsed < 200*time.Millisecond,
		"Event should have been delivered immediately, took %s", elapsed)
	require.Equal(t, index, got.Index)

	// Schedule an event publish in a while
	index++
	start = time.Now()
	time.AfterFunc(200*time.Millisecond, func() {
		publishTestEvent(index, eb, "test")
	})

	// Next call should block until event is delivered
	got, err = sub.Next(ctx)
	elapsed = time.Since(start)
	require.NoError(t, err)
	require.True(t, elapsed > 200*time.Millisecond,
		"Event should have been delivered after blocking 200ms, took %s", elapsed)
	require.True(t, elapsed < 2*time.Second,
		"Event should have been delivered after short time, took %s", elapsed)
	require.Equal(t, index, got.Index)

	// Event with wrong key should not be delivered. Deliver a good message right
	// so we don't have to block test thread forever or cancel func yet.
	index++
	publishTestEvent(index, eb, "nope")
	index++
	publishTestEvent(index, eb, "test")

	start = time.Now()
	got, err = sub.Next(ctx)
	elapsed = time.Since(start)
	require.NoError(t, err)
	require.True(t, elapsed < 200*time.Millisecond,
		"Event should have been delivered immediately, took %s", elapsed)
	require.Equal(t, index, got.Index)
	require.Equal(t, "test", got.Payload.(simplePayload).key)

	// Cancelling the subscription context should unblock Next
	start = time.Now()
	time.AfterFunc(200*time.Millisecond, func() {
		cancel()
	})

	_, err = sub.Next(ctx)
	elapsed = time.Since(start)
	require.Error(t, err)
	require.True(t, elapsed > 200*time.Millisecond,
		"Event should have been delivered after blocking 200ms, took %s", elapsed)
	require.True(t, elapsed < 2*time.Second,
		"Event should have been delivered after short time, took %s", elapsed)
}

func TestSubscription_Close(t *testing.T) {
	eb := newEventBuffer()
	index := uint64(100)
	startHead := eb.Head()

	// Start with an event in the buffer
	publishTestEvent(index, eb, "test")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	req := SubscribeRequest{
		Topic: testTopic,
		Key:   "test",
	}
	sub := newSubscription(req, startHead, noopUnSub)

	// First call to sub.Next should return our published event immediately
	start := time.Now()
	got, err := sub.Next(ctx)
	elapsed := time.Since(start)
	require.NoError(t, err)
	require.True(t, elapsed < 200*time.Millisecond,
		"Event should have been delivered immediately, took %s", elapsed)
	require.Equal(t, index, got.Index)

	// Schedule a Close simulating the server deciding this subscroption
	// needs to reset (e.g. on ACL perm change).
	start = time.Now()
	time.AfterFunc(200*time.Millisecond, func() {
		sub.forceClose()
	})

	_, err = sub.Next(ctx)
	elapsed = time.Since(start)
	require.Error(t, err)
	require.Equal(t, ErrSubForceClosed, err)
	require.True(t, elapsed > 200*time.Millisecond,
		"Reload should have happened after blocking 200ms, took %s", elapsed)
	require.True(t, elapsed < 2*time.Second,
		"Reload should have been delivered after short time, took %s", elapsed)
}

func publishTestEvent(index uint64, b *eventBuffer, key string) {
	e := Event{
		Index:   index,
		Topic:   testTopic,
		Payload: simplePayload{key: key},
	}
	b.Append([]Event{e})
}

func TestFilter_NoKey(t *testing.T) {
	events := make(PayloadEvents, 0, 5)
	events = append(events, newSimpleEvent("One", 102), newSimpleEvent("Two", 102))

	req := SubscribeRequest{Topic: testTopic}
	actual, ok := filterByKey(req, events)
	require.True(t, ok)
	require.Equal(t, Event{Topic: testTopic, Index: 102, Payload: events}, actual)

	// test that a new array was not allocated
	require.Equal(t, cap(actual.Payload.(PayloadEvents)), 5)
}

func newSimpleEvent(key string, index uint64) Event {
	return Event{Index: index, Payload: simplePayload{key: key}}
}

func TestFilter_WithKey_AllEventsMatch(t *testing.T) {
	events := make(PayloadEvents, 0, 5)
	events = append(events, newSimpleEvent("Same", 103), newSimpleEvent("Same", 103))

	req := SubscribeRequest{Topic: testTopic, Key: "Same"}
	actual, ok := filterByKey(req, events)
	require.True(t, ok)
	expected := Event{Topic: testTopic, Index: 103, Payload: events}
	require.Equal(t, expected, actual)

	// test that a new array was not allocated
	require.Equal(t, 5, cap(actual.Payload.(PayloadEvents)))
}

func TestFilter_WithKey_SomeEventsMatch(t *testing.T) {
	events := make([]Event, 0, 5)
	events = append(events,
		newSimpleEvent("Same", 104),
		newSimpleEvent("Other", 0),
		newSimpleEvent("Same", 0))

	req := SubscribeRequest{Topic: testTopic, Key: "Same"}
	actual, ok := filterByKey(req, events)
	require.True(t, ok)
	expected := Event{
		Topic:   testTopic,
		Index:   104,
		Payload: PayloadEvents{newSimpleEvent("Same", 104), newSimpleEvent("Same", 0)},
	}
	require.Equal(t, expected, actual)

	// test that a new array was allocated with the correct size
	require.Equal(t, cap(actual.Payload.(PayloadEvents)), 2)
}

func TestFilter_WithKey_NoEventsMatch(t *testing.T) {
	events := make([]Event, 0, 5)
	events = append(events, newSimpleEvent("Same", 0), newSimpleEvent("Same", 0))

	req := SubscribeRequest{Topic: testTopic, Key: "Other"}
	_, ok := filterByKey(req, events)
	require.False(t, ok)
}
<commit_msg>stream: Move FilterByKey events to a table<commit_after>package stream

import (
	"context"
	"testing"
	time "time"

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

func noopUnSub() {}

func TestSubscription(t *testing.T) {
	eb := newEventBuffer()

	index := uint64(100)

	startHead := eb.Head()

	// Start with an event in the buffer
	publishTestEvent(index, eb, "test")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	req := SubscribeRequest{
		Topic: testTopic,
		Key:   "test",
	}
	sub := newSubscription(req, startHead, noopUnSub)

	// First call to sub.Next should return our published event immediately
	start := time.Now()
	got, err := sub.Next(ctx)
	elapsed := time.Since(start)
	require.NoError(t, err)
	require.True(t, elapsed < 200*time.Millisecond,
		"Event should have been delivered immediately, took %s", elapsed)
	require.Equal(t, index, got.Index)

	// Schedule an event publish in a while
	index++
	start = time.Now()
	time.AfterFunc(200*time.Millisecond, func() {
		publishTestEvent(index, eb, "test")
	})

	// Next call should block until event is delivered
	got, err = sub.Next(ctx)
	elapsed = time.Since(start)
	require.NoError(t, err)
	require.True(t, elapsed > 200*time.Millisecond,
		"Event should have been delivered after blocking 200ms, took %s", elapsed)
	require.True(t, elapsed < 2*time.Second,
		"Event should have been delivered after short time, took %s", elapsed)
	require.Equal(t, index, got.Index)

	// Event with wrong key should not be delivered. Deliver a good message right
	// so we don't have to block test thread forever or cancel func yet.
	index++
	publishTestEvent(index, eb, "nope")
	index++
	publishTestEvent(index, eb, "test")

	start = time.Now()
	got, err = sub.Next(ctx)
	elapsed = time.Since(start)
	require.NoError(t, err)
	require.True(t, elapsed < 200*time.Millisecond,
		"Event should have been delivered immediately, took %s", elapsed)
	require.Equal(t, index, got.Index)
	require.Equal(t, "test", got.Payload.(simplePayload).key)

	// Cancelling the subscription context should unblock Next
	start = time.Now()
	time.AfterFunc(200*time.Millisecond, func() {
		cancel()
	})

	_, err = sub.Next(ctx)
	elapsed = time.Since(start)
	require.Error(t, err)
	require.True(t, elapsed > 200*time.Millisecond,
		"Event should have been delivered after blocking 200ms, took %s", elapsed)
	require.True(t, elapsed < 2*time.Second,
		"Event should have been delivered after short time, took %s", elapsed)
}

func TestSubscription_Close(t *testing.T) {
	eb := newEventBuffer()
	index := uint64(100)
	startHead := eb.Head()

	// Start with an event in the buffer
	publishTestEvent(index, eb, "test")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	req := SubscribeRequest{
		Topic: testTopic,
		Key:   "test",
	}
	sub := newSubscription(req, startHead, noopUnSub)

	// First call to sub.Next should return our published event immediately
	start := time.Now()
	got, err := sub.Next(ctx)
	elapsed := time.Since(start)
	require.NoError(t, err)
	require.True(t, elapsed < 200*time.Millisecond,
		"Event should have been delivered immediately, took %s", elapsed)
	require.Equal(t, index, got.Index)

	// Schedule a Close simulating the server deciding this subscroption
	// needs to reset (e.g. on ACL perm change).
	start = time.Now()
	time.AfterFunc(200*time.Millisecond, func() {
		sub.forceClose()
	})

	_, err = sub.Next(ctx)
	elapsed = time.Since(start)
	require.Error(t, err)
	require.Equal(t, ErrSubForceClosed, err)
	require.True(t, elapsed > 200*time.Millisecond,
		"Reload should have happened after blocking 200ms, took %s", elapsed)
	require.True(t, elapsed < 2*time.Second,
		"Reload should have been delivered after short time, took %s", elapsed)
}

func publishTestEvent(index uint64, b *eventBuffer, key string) {
	e := Event{
		Index:   index,
		Topic:   testTopic,
		Payload: simplePayload{key: key},
	}
	b.Append([]Event{e})
}

func newSimpleEvent(key string, index uint64) Event {
	return Event{Index: index, Payload: simplePayload{key: key}}
}

func TestFilterByKey(t *testing.T) {
	type testCase struct {
		name        string
		req         SubscribeRequest
		events      []Event
		expectEvent bool
		expected    Event
		expectedCap int
	}

	fn := func(t *testing.T, tc testCase) {
		events := make(PayloadEvents, 0, 5)
		events = append(events, tc.events...)

		actual, ok := filterByKey(tc.req, events)
		require.Equal(t, tc.expectEvent, ok)
		if !tc.expectEvent {
			return
		}

		require.Equal(t, tc.expected, actual)
		// test if there was a new array allocated or not
		require.Equal(t, tc.expectedCap, cap(actual.Payload.(PayloadEvents)))
	}

	var testCases = []testCase{
		{
			name: "all events match, no key or namespace",
			req:  SubscribeRequest{Topic: testTopic},
			events: []Event{
				newSimpleEvent("One", 102),
				newSimpleEvent("Two", 102)},
			expectEvent: true,
			expected: Event{
				Topic: testTopic,
				Index: 102,
				Payload: PayloadEvents{
					newSimpleEvent("One", 102),
					newSimpleEvent("Two", 102)}},
			expectedCap: 5,
		},
		{
			name: "all events match, no namespace",
			req:  SubscribeRequest{Topic: testTopic, Key: "Same"},
			events: []Event{
				newSimpleEvent("Same", 103),
				newSimpleEvent("Same", 103)},
			expectEvent: true,
			expected: Event{
				Topic: testTopic,
				Index: 103,
				Payload: PayloadEvents{
					newSimpleEvent("Same", 103),
					newSimpleEvent("Same", 103)}},
			expectedCap: 5,
		},
		// TODO: all events match, no key
		{
			name: "some evens match, no namespace",
			req:  SubscribeRequest{Topic: testTopic, Key: "Same"},
			events: []Event{
				newSimpleEvent("Same", 104),
				newSimpleEvent("Other", 104),
				newSimpleEvent("Same", 104)},
			expectEvent: true,
			expected: Event{
				Topic: testTopic,
				Index: 104,
				Payload: PayloadEvents{
					newSimpleEvent("Same", 104),
					newSimpleEvent("Same", 104)}},
			expectedCap: 2,
		},
		// TODO: some events match, no key
		{
			name: "no events match key",
			req:  SubscribeRequest{Topic: testTopic, Key: "Other"},
			events: []Event{
				newSimpleEvent("Same", 0),
				newSimpleEvent("Same", 0)},
		},
		// TODO: no events match namespace
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			fn(t, tc)
		})
	}
}
<|endoftext|>


<commit_before>package healthcheck_test

import (
	"fmt"
	"net"
	"net/http"
	"runtime"
	"time"

	"code.cloudfoundry.org/healthcheck"

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

var _ = Describe("HealthCheck", func() {
	itReturnsHealthCheckError := func(healthCheck func() error, code int, reason string) {
		err := healthCheck()
		Expect(err).To(HaveOccurred())
		Expect(err).To(BeAssignableToTypeOf(healthcheck.HealthCheckError{}))
		hErr := err.(healthcheck.HealthCheckError)
		Expect(hErr.Code).To(Equal(code))
		Expect(hErr.Message).To(ContainSubstring(reason))
	}

	var (
		server     *ghttp.Server
		serverAddr string

		ip string

		uri         string
		port        string
		timeout     time.Duration
		serverDelay time.Duration

		hc      healthcheck.HealthCheck
		handler func(http.ResponseWriter, *http.Request)
	)

	BeforeEach(func() {
		ip = getNonLoopbackIP()
		server = ghttp.NewUnstartedServer()

		listener, err := net.Listen("tcp", ip+":0")
		Expect(err).NotTo(HaveOccurred())

		timeout = 100 * time.Millisecond
		serverDelay = 0

		server.HTTPTestServer.Listener = listener
		serverAddr = listener.Addr().String()
		_, port, err = net.SplitHostPort(serverAddr)
		Expect(err).NotTo(HaveOccurred())

		handler = func(http.ResponseWriter, *http.Request) {
			time.Sleep(serverDelay)
		}

	})

	JustBeforeEach(func() {

		server.RouteToHandler("GET", "/api/_ping", handler)
		server.Start()

		hc = healthcheck.NewHealthCheck("tcp", uri, port, timeout)
	})

	AfterEach(func() {
		if server != nil {
			server.CloseClientConnections()
			server.Close()
		}
	})

	Describe("check interfaces", func() {
		It("succeeds when there are healthy interfaces", func() {
			interfaces, err := net.Interfaces()
			Expect(err).NotTo(HaveOccurred())

			err = hc.CheckInterfaces(interfaces)
			Expect(err).NotTo(HaveOccurred())
		})

		Context("when the server is failing", func() {
			AfterEach(func() {
				server = nil
			})

			It("fails appropriately when there are unhealthy interfaces", func() {
				server.Close()

				interfaces, err := net.Interfaces()
				Expect(err).NotTo(HaveOccurred())

				err = hc.CheckInterfaces(interfaces)
				Expect(err).To(HaveOccurred())
				Expect(err).To(BeAssignableToTypeOf(healthcheck.HealthCheckError{}))

				hErr := err.(healthcheck.HealthCheckError)
				// fails with different error codes on Linux (4) or OSX (64)
				// check to see it was not the NO interfaces error (3)
				Expect(hErr.Code).ToNot(Equal(3))
			})
		})

		It("fails appropriately when there are no interfaces", func() {
			err := hc.CheckInterfaces(nil)
			Expect(err).To(HaveOccurred())
			Expect(err).To(BeAssignableToTypeOf(healthcheck.HealthCheckError{}))

			hErr := err.(healthcheck.HealthCheckError)
			Expect(hErr.Code).To(Equal(3))
			Expect(hErr.Message).To(ContainSubstring("failure to find suitable interface"))
		})
	})

	Describe("port healthcheck", func() {
		portHealthCheck := func() error {
			return hc.PortHealthCheck(ip)
		}

		BeforeEach(func() {
			uri = ""
		})

		Context("when the address is listening", func() {
			It("succeeds", func() {
				Expect(portHealthCheck()).To(Succeed())
			})
		})

		Context("when the address is not listening", func() {
			BeforeEach(func() {
				port = "-1"
			})

			It("returns healthcheck error with code 4 with an appropriate message", func() {
				errMsg := fmt.Sprintf(
					"Failed to make TCP connection to port %s: connection refused",
					port,
				)
				itReturnsHealthCheckError(portHealthCheck, 4, errMsg)
			})
		})

		Context("when the server is slow in responding", func() {
			BeforeEach(func() {
				if runtime.GOOS == "windows" {
					Skip("TCP timeout tests are unreliable on Windows")
				}

				timeout = time.Nanosecond
			})

			It("returns healthcheck error with code 64 with an appropriate message", func() {
				errMsg := fmt.Sprintf(
					"Failed to make TCP connection to port %s: timed out after %.2f seconds",
					port,
					timeout.Seconds(),
				)
				itReturnsHealthCheckError(portHealthCheck, 64, errMsg)
			})
		})
	})

	Describe("http healthcheck", func() {
		httpHealthCheck := func() error {
			return hc.HTTPHealthCheck(ip)
		}

		BeforeEach(func() {
			uri = "/api/_ping"
		})

		Context("when the healthcheck is properly invoked", func() {
			Context("when the address is listening", func() {
				It("succeeds", func() {
					Expect(httpHealthCheck()).To(Succeed())
				})
			})

			Context("when the address returns error http code", func() {
				JustBeforeEach(func() {
					server.RouteToHandler("GET", "/api/_ping", ghttp.RespondWith(500, ""))
				})

				It("returns healthcheck error with code 6 with an appropriate message", func() {
					errMsg := fmt.Sprintf(
						"Failed to make HTTP request to '%s' on port %s: received status code 500 in",
						uri,
						port,
					)
					itReturnsHealthCheckError(httpHealthCheck, 6, errMsg)
				})
			})

			Context("when the address is not listening", func() {
				BeforeEach(func() {
					port = "9876"
				})

				It("returns healthcheck error with code 5 with an appropriate message", func() {
					errMsg := fmt.Sprintf(
						"Failed to make HTTP request to '%s' on port %s: connection refused",
						uri,
						port,
					)
					itReturnsHealthCheckError(httpHealthCheck, 5, errMsg)
				})
			})

			Context("when the server is too slow to respond", func() {
				BeforeEach(func() {
					timeout = time.Nanosecond
					serverDelay = time.Second
				})

				It("returns healthcheck error with code 65 with an appropriate message", func() {
					errMsg := fmt.Sprintf(
						"Failed to make HTTP request to '%s' on port %s: timed out after %.2f seconds",
						uri,
						port,
						timeout.Seconds(),
					)
					itReturnsHealthCheckError(httpHealthCheck, 65, errMsg)
				})
			})

			Context("with a tls-aware endpoint", func() {
				var request *http.Request
				BeforeEach(func() {
					uri = "/api/_ping"
					handler = func(resp http.ResponseWriter, req *http.Request) {
						request = req
						time.Sleep(serverDelay)
					}
				})

				It("should set the x-Forwarded-proto header", func() {
					err := hc.HTTPHealthCheck(ip)
					Expect(err).NotTo(HaveOccurred())
					Expect(request.Header.Get("X-Forwarded-Proto")).To(Equal("https"))
				})
			})

		})
	})
})

func getNonLoopbackIP() string {
	interfaces, err := net.Interfaces()
	Expect(err).NotTo(HaveOccurred())
	for _, intf := range interfaces {
		addrs, err := intf.Addrs()
		if err != nil {
			continue
		}

		for _, a := range addrs {
			if ipnet, ok := a.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
				if ipnet.IP.To4() != nil {
					return ipnet.IP.String()
				}
			}
		}
	}
	Fail("no non-loopback address found")
	panic("non-reachable")
}
<commit_msg>error on windows, try giving an out-of-range port<commit_after>package healthcheck_test

import (
	"fmt"
	"net"
	"net/http"
	"runtime"
	"time"

	"code.cloudfoundry.org/healthcheck"

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

var _ = Describe("HealthCheck", func() {
	itReturnsHealthCheckError := func(healthCheck func() error, code int, reason string) {
		err := healthCheck()
		Expect(err).To(HaveOccurred())
		Expect(err).To(BeAssignableToTypeOf(healthcheck.HealthCheckError{}))
		hErr := err.(healthcheck.HealthCheckError)
		Expect(hErr.Code).To(Equal(code))
		Expect(hErr.Message).To(ContainSubstring(reason))
	}

	var (
		server     *ghttp.Server
		serverAddr string

		ip string

		uri         string
		port        string
		timeout     time.Duration
		serverDelay time.Duration

		hc      healthcheck.HealthCheck
		handler func(http.ResponseWriter, *http.Request)
	)

	BeforeEach(func() {
		ip = getNonLoopbackIP()
		server = ghttp.NewUnstartedServer()

		listener, err := net.Listen("tcp", ip+":0")
		Expect(err).NotTo(HaveOccurred())

		timeout = 100 * time.Millisecond
		serverDelay = 0

		server.HTTPTestServer.Listener = listener
		serverAddr = listener.Addr().String()
		_, port, err = net.SplitHostPort(serverAddr)
		Expect(err).NotTo(HaveOccurred())

		handler = func(http.ResponseWriter, *http.Request) {
			time.Sleep(serverDelay)
		}

	})

	JustBeforeEach(func() {

		server.RouteToHandler("GET", "/api/_ping", handler)
		server.Start()

		hc = healthcheck.NewHealthCheck("tcp", uri, port, timeout)
	})

	AfterEach(func() {
		if server != nil {
			server.CloseClientConnections()
			server.Close()
		}
	})

	Describe("check interfaces", func() {
		It("succeeds when there are healthy interfaces", func() {
			interfaces, err := net.Interfaces()
			Expect(err).NotTo(HaveOccurred())

			err = hc.CheckInterfaces(interfaces)
			Expect(err).NotTo(HaveOccurred())
		})

		Context("when the server is failing", func() {
			AfterEach(func() {
				server = nil
			})

			It("fails appropriately when there are unhealthy interfaces", func() {
				server.Close()

				interfaces, err := net.Interfaces()
				Expect(err).NotTo(HaveOccurred())

				err = hc.CheckInterfaces(interfaces)
				Expect(err).To(HaveOccurred())
				Expect(err).To(BeAssignableToTypeOf(healthcheck.HealthCheckError{}))

				hErr := err.(healthcheck.HealthCheckError)
				// fails with different error codes on Linux (4) or OSX (64)
				// check to see it was not the NO interfaces error (3)
				Expect(hErr.Code).ToNot(Equal(3))
			})
		})

		It("fails appropriately when there are no interfaces", func() {
			err := hc.CheckInterfaces(nil)
			Expect(err).To(HaveOccurred())
			Expect(err).To(BeAssignableToTypeOf(healthcheck.HealthCheckError{}))

			hErr := err.(healthcheck.HealthCheckError)
			Expect(hErr.Code).To(Equal(3))
			Expect(hErr.Message).To(ContainSubstring("failure to find suitable interface"))
		})
	})

	Describe("port healthcheck", func() {
		portHealthCheck := func() error {
			return hc.PortHealthCheck(ip)
		}

		BeforeEach(func() {
			uri = ""
		})

		Context("when the address is listening", func() {
			It("succeeds", func() {
				Expect(portHealthCheck()).To(Succeed())
			})
		})

		Context("when the address is not listening", func() {
			BeforeEach(func() {
				port = "-1"
			})

			It("returns healthcheck error with code 4 with an appropriate message", func() {
				errMsg := fmt.Sprintf(
					"Failed to make TCP connection to port %s: connection refused",
					port,
				)
				itReturnsHealthCheckError(portHealthCheck, 4, errMsg)
			})
		})

		Context("when the server is slow in responding", func() {
			BeforeEach(func() {
				if runtime.GOOS == "windows" {
					Skip("TCP timeout tests are unreliable on Windows")
				}

				timeout = time.Nanosecond
			})

			It("returns healthcheck error with code 64 with an appropriate message", func() {
				errMsg := fmt.Sprintf(
					"Failed to make TCP connection to port %s: timed out after %.2f seconds",
					port,
					timeout.Seconds(),
				)
				itReturnsHealthCheckError(portHealthCheck, 64, errMsg)
			})
		})
	})

	Describe("http healthcheck", func() {
		httpHealthCheck := func() error {
			return hc.HTTPHealthCheck(ip)
		}

		BeforeEach(func() {
			uri = "/api/_ping"
		})

		Context("when the healthcheck is properly invoked", func() {
			Context("when the address is listening", func() {
				It("succeeds", func() {
					Expect(httpHealthCheck()).To(Succeed())
				})
			})

			Context("when the address returns error http code", func() {
				JustBeforeEach(func() {
					server.RouteToHandler("GET", "/api/_ping", ghttp.RespondWith(500, ""))
				})

				It("returns healthcheck error with code 6 with an appropriate message", func() {
					errMsg := fmt.Sprintf(
						"Failed to make HTTP request to '%s' on port %s: received status code 500 in",
						uri,
						port,
					)
					itReturnsHealthCheckError(httpHealthCheck, 6, errMsg)
				})
			})

			Context("when the address is not listening", func() {
				BeforeEach(func() {
					port = "100003"
				})

				It("returns healthcheck error with code 5 with an appropriate message", func() {
					errMsg := fmt.Sprintf(
						"Failed to make HTTP request to '%s' on port %s: connection refused",
						uri,
						port,
					)
					itReturnsHealthCheckError(httpHealthCheck, 5, errMsg)
				})
			})

			Context("when the server is too slow to respond", func() {
				BeforeEach(func() {
					timeout = time.Nanosecond
					serverDelay = time.Second
				})

				It("returns healthcheck error with code 65 with an appropriate message", func() {
					errMsg := fmt.Sprintf(
						"Failed to make HTTP request to '%s' on port %s: timed out after %.2f seconds",
						uri,
						port,
						timeout.Seconds(),
					)
					itReturnsHealthCheckError(httpHealthCheck, 65, errMsg)
				})
			})

			Context("with a tls-aware endpoint", func() {
				var request *http.Request
				BeforeEach(func() {
					uri = "/api/_ping"
					handler = func(resp http.ResponseWriter, req *http.Request) {
						request = req
						time.Sleep(serverDelay)
					}
				})

				It("should set the x-Forwarded-proto header", func() {
					err := hc.HTTPHealthCheck(ip)
					Expect(err).NotTo(HaveOccurred())
					Expect(request.Header.Get("X-Forwarded-Proto")).To(Equal("https"))
				})
			})

		})
	})
})

func getNonLoopbackIP() string {
	interfaces, err := net.Interfaces()
	Expect(err).NotTo(HaveOccurred())
	for _, intf := range interfaces {
		addrs, err := intf.Addrs()
		if err != nil {
			continue
		}

		for _, a := range addrs {
			if ipnet, ok := a.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
				if ipnet.IP.To4() != nil {
					return ipnet.IP.String()
				}
			}
		}
	}
	Fail("no non-loopback address found")
	panic("non-reachable")
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

func TestMakeMovePawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(0, 2))
	from := byte(CartesianToIndex(0, 1))
	move := Move{from, to, MoveQuiet, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) != PAWN {
		t.Fail()
	}
}

func TestMakeMoveWhitePawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(0, 3))
	from := byte(CartesianToIndex(0, 1))
	move := Move{from, to, MoveDoublePush, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || board.EnPassant != CartesianToIndex(0, 2) {
		t.Fail()
	}
}

func TestMakeMoveBlackPawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	to := byte(CartesianToIndex(0, 4))
	from := byte(CartesianToIndex(0, 6))
	move := Move{from, to, MoveDoublePush, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || board.EnPassant != CartesianToIndex(0, 5) {
		t.Fail()
	}
}

func TestMakeMoveWhitePawnCapture(t *testing.T) {
	board, err := Parse("8/8/8/3p4/4P3/8/8/8 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(3, 4))
	from := byte(CartesianToIndex(4, 3))
	move := Move{from, to, MoveCapture, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || GetSide(board.Data[to]) != WHITE {
		t.Fail()
	}
}

func TestMakeMoveBlackPawnCapture(t *testing.T) {
	board, err := Parse("8/8/8/3p4/4P3/8/8/8 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(4, 3))
	from := byte(CartesianToIndex(3, 4))
	move := Move{from, to, MoveCapture, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || GetSide(board.Data[to]) != BLACK {
		t.Fail()
	}
}

func IsMoveInMoveList(t *testing.T, list []Move, from, to, kind byte) bool {
	for _, m := range list {
		if m.To == to && m.From == from && m.Kind == kind {
			return true
		}
	}
	return false
}

func TestMoveGenWhitePawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 1)), byte(CartesianToIndex(0, 2)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 6)), byte(CartesianToIndex(0, 5)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenWhitePawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 1)), byte(CartesianToIndex(0, 3)), MoveDoublePush) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 6)), byte(CartesianToIndex(0, 4)), MoveDoublePush) {
		t.Fail()
	}
}

func TestMoveGenSlider(t *testing.T) {
	board, err := Parse("4k3/8/8/1r6/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(1, 4)), byte(CartesianToIndex(1, 7)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenNonSlider(t *testing.T) {
	board, err := Parse("4k3/8/8/1r6/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 7)), byte(CartesianToIndex(4, 6)), MoveQuiet) {
		t.Fail()
	}
	if IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 7)), byte(CartesianToIndex(4, 3)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenWhitePawnCap(t *testing.T) {
	board, err := Parse("4k3/8/8/3p4/4P3/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 3)), byte(CartesianToIndex(3, 4)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnCap(t *testing.T) {
	board, err := Parse("4k3/8/8/3p4/4P3/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(3, 4)), byte(CartesianToIndex(4, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenSliderCap(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1r1P4/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(1, 3)), byte(CartesianToIndex(3, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenNonSliderCap(t *testing.T) {
	board, err := Parse("4k3/8/2n5/8/3P4/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	t.Log(moves)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(2, 5)), byte(CartesianToIndex(3, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestPerft(t *testing.T) {
	tperft(t, 1, 20)
	tperft(t, 2, 400)
	tperft(t, 3, 8902)
	tperft(t, 4, 197281)
	tperft(t, 5, 4865609)
}

func tperft(t *testing.T, depth int, expected uint64) {
	result := DoDividePerft(depth)
	t.Log("Result of perft(", depth, "): ", result)
	if result != expected {
		t.FailNow()
	}
}

func TestSquareNotAttacked(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/3p4/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if squareattacked(board, CartesianToIndex(4, 4), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByWhitePawn(t *testing.T) {
	board, err := Parse("4k3/3P4/8/8/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 7), WHITE) {
		t.FailNow()
	}
}

func TestSquareAttackedByBlackPawn(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/3p4/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByQueenFile(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/4q3/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByQueenDiag(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1q6/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByRook(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/4r3/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByBishop(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1b6/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByKnight(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/2n5/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestParseMoveParsesValidMove(t *testing.T) {
	_, err := ParseMove("e2e4")
	if err != nil {
		t.Fail()
	}
}

func TestParseMoveRejectsTooLongMove(t *testing.T) {
	_, err := ParseMove("e2e4e7e8f6")
	if err == nil {
		t.Fail()
	}
}

func TestParseMoveRejectsInvalidFrom(t *testing.T) {
	_, err := ParseMove("l9e4")
	if err == nil {
		t.Fail()
	}
}

func TestParseMoveRejectsInvalidTo(t *testing.T) {
	_, err := ParseMove("e2l9")
	if err == nil {
		t.Fail()
	}
}

func TestMakeMoveEnPassantWhite(t *testing.T) {
	board, err := Parse("rnbqkbnr/1pp1pppp/p7/3pP3/8/8/PPPP1PPP/RNBQKBNR w KQkq d6 0 1")
	if err != nil {
		t.FailNow()
	}
	from, _ := AlgebraicToIndex("e5")
	to, _ := AlgebraicToIndex("d6")
	taken, _ := AlgebraicToIndex("d5")
	MakeMove(board, &Move{from, to, MoveEnPassant, EMPTY, 0})
	if GetPiece(board.Data[taken]) != EMPTY || GetPiece(board.Data[to]) != PAWN {
		t.Fail()
	}
}

func TestMakeMoveEnPassantBlack(t *testing.T) {
	board, err := Parse("rnbqkbnr/ppp1pppp/8/8/3pP3/PP6/2PP1PPP/RNBQKBNR b KQkq e3 0 1")
	if err != nil {
		t.FailNow()
	}
	from, _ := AlgebraicToIndex("d4")
	to, _ := AlgebraicToIndex("e3")
	taken, _ := AlgebraicToIndex("e4")
	MakeMove(board, &Move{from, to, MoveEnPassant, EMPTY, 0})
	if GetPiece(board.Data[taken]) != EMPTY || GetPiece(board.Data[to]) != PAWN {
		t.Fail()
	}
}

func TestMoveGenEnPassant(t *testing.T) {
	board, err := Parse("rnbqkbnr/1pp1pppp/p7/3pP3/8/8/PPPP1PPP/RNBQKBNR w KQkq d6 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	from, _ := AlgebraicToIndex("e5")
	to, _ := AlgebraicToIndex("d6")
	if !IsMoveInMoveList(t, moves, from, to, MoveEnPassant) {
		t.Fail()
	}
}
<commit_msg>cover unmaking black enpassant<commit_after>package main

import (
	"testing"
)

func TestMakeMovePawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(0, 2))
	from := byte(CartesianToIndex(0, 1))
	move := Move{from, to, MoveQuiet, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) != PAWN {
		t.Fail()
	}
}

func TestMakeMoveWhitePawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(0, 3))
	from := byte(CartesianToIndex(0, 1))
	move := Move{from, to, MoveDoublePush, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || board.EnPassant != CartesianToIndex(0, 2) {
		t.Fail()
	}
}

func TestMakeMoveBlackPawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	to := byte(CartesianToIndex(0, 4))
	from := byte(CartesianToIndex(0, 6))
	move := Move{from, to, MoveDoublePush, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || board.EnPassant != CartesianToIndex(0, 5) {
		t.Fail()
	}
}

func TestMakeMoveWhitePawnCapture(t *testing.T) {
	board, err := Parse("8/8/8/3p4/4P3/8/8/8 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(3, 4))
	from := byte(CartesianToIndex(4, 3))
	move := Move{from, to, MoveCapture, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || GetSide(board.Data[to]) != WHITE {
		t.Fail()
	}
}

func TestMakeMoveBlackPawnCapture(t *testing.T) {
	board, err := Parse("8/8/8/3p4/4P3/8/8/8 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(4, 3))
	from := byte(CartesianToIndex(3, 4))
	move := Move{from, to, MoveCapture, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || GetSide(board.Data[to]) != BLACK {
		t.Fail()
	}
}

func IsMoveInMoveList(t *testing.T, list []Move, from, to, kind byte) bool {
	for _, m := range list {
		if m.To == to && m.From == from && m.Kind == kind {
			return true
		}
	}
	return false
}

func TestMoveGenWhitePawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 1)), byte(CartesianToIndex(0, 2)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 6)), byte(CartesianToIndex(0, 5)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenWhitePawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 1)), byte(CartesianToIndex(0, 3)), MoveDoublePush) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 6)), byte(CartesianToIndex(0, 4)), MoveDoublePush) {
		t.Fail()
	}
}

func TestMoveGenSlider(t *testing.T) {
	board, err := Parse("4k3/8/8/1r6/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(1, 4)), byte(CartesianToIndex(1, 7)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenNonSlider(t *testing.T) {
	board, err := Parse("4k3/8/8/1r6/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 7)), byte(CartesianToIndex(4, 6)), MoveQuiet) {
		t.Fail()
	}
	if IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 7)), byte(CartesianToIndex(4, 3)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenWhitePawnCap(t *testing.T) {
	board, err := Parse("4k3/8/8/3p4/4P3/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 3)), byte(CartesianToIndex(3, 4)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnCap(t *testing.T) {
	board, err := Parse("4k3/8/8/3p4/4P3/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(3, 4)), byte(CartesianToIndex(4, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenSliderCap(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1r1P4/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(1, 3)), byte(CartesianToIndex(3, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenNonSliderCap(t *testing.T) {
	board, err := Parse("4k3/8/2n5/8/3P4/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	t.Log(moves)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(2, 5)), byte(CartesianToIndex(3, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestPerft(t *testing.T) {
	tperft(t, 1, 20)
	tperft(t, 2, 400)
	tperft(t, 3, 8902)
	tperft(t, 4, 197281)
	tperft(t, 5, 4865609)
}

func tperft(t *testing.T, depth int, expected uint64) {
	result := DoDividePerft(depth)
	t.Log("Result of perft(", depth, "): ", result)
	if result != expected {
		t.FailNow()
	}
}

func TestSquareNotAttacked(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/3p4/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if squareattacked(board, CartesianToIndex(4, 4), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByWhitePawn(t *testing.T) {
	board, err := Parse("4k3/3P4/8/8/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 7), WHITE) {
		t.FailNow()
	}
}

func TestSquareAttackedByBlackPawn(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/3p4/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByQueenFile(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/4q3/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByQueenDiag(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1q6/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByRook(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/4r3/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByBishop(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1b6/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByKnight(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/2n5/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestParseMoveParsesValidMove(t *testing.T) {
	_, err := ParseMove("e2e4")
	if err != nil {
		t.Fail()
	}
}

func TestParseMoveRejectsTooLongMove(t *testing.T) {
	_, err := ParseMove("e2e4e7e8f6")
	if err == nil {
		t.Fail()
	}
}

func TestParseMoveRejectsInvalidFrom(t *testing.T) {
	_, err := ParseMove("l9e4")
	if err == nil {
		t.Fail()
	}
}

func TestParseMoveRejectsInvalidTo(t *testing.T) {
	_, err := ParseMove("e2l9")
	if err == nil {
		t.Fail()
	}
}

func TestMakeMoveEnPassantWhite(t *testing.T) {
	board, err := Parse("rnbqkbnr/1pp1pppp/p7/3pP3/8/8/PPPP1PPP/RNBQKBNR w KQkq d6 0 1")
	if err != nil {
		t.FailNow()
	}
	from, _ := AlgebraicToIndex("e5")
	to, _ := AlgebraicToIndex("d6")
	taken, _ := AlgebraicToIndex("d5")
	MakeMove(board, &Move{from, to, MoveEnPassant, EMPTY, 0})
	if GetPiece(board.Data[taken]) != EMPTY || GetPiece(board.Data[to]) != PAWN {
		t.Fail()
	}
}

func TestMakeMoveEnPassantBlack(t *testing.T) {
	board, err := Parse("rnbqkbnr/ppp1pppp/8/8/3pP3/PP6/2PP1PPP/RNBQKBNR b KQkq e3 0 1")
	if err != nil {
		t.FailNow()
	}
	from, _ := AlgebraicToIndex("d4")
	to, _ := AlgebraicToIndex("e3")
	taken, _ := AlgebraicToIndex("e4")
	MakeMove(board, &Move{from, to, MoveEnPassant, EMPTY, 0})
	if GetPiece(board.Data[taken]) != EMPTY || GetPiece(board.Data[to]) != PAWN {
		t.Fail()
	}
}

func TestMoveGenEnPassant(t *testing.T) {
	board, err := Parse("rnbqkbnr/1pp1pppp/p7/3pP3/8/8/PPPP1PPP/RNBQKBNR w KQkq d6 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	from, _ := AlgebraicToIndex("e5")
	to, _ := AlgebraicToIndex("d6")
	if !IsMoveInMoveList(t, moves, from, to, MoveEnPassant) {
		t.Fail()
	}
}

func TestUnmakeMoveBlackEnPassant(t *testing.T) {
	board, err := Parse("rnbqkbnr/ppp1pppp/8/8/3pP3/PP6/2PP1PPP/RNBQKBNR b KQkq e3 0 1")
	if err != nil {
		t.FailNow()
	}
	from, _ := AlgebraicToIndex("d4")
	to, _ := AlgebraicToIndex("e3")
	taken, _ := AlgebraicToIndex("e4")
	move := &Move{from, to, MoveEnPassant, EMPTY, 0}
	undo := MakeMove(board, move)
	UnmakeMove(board, move, undo)
	if GetPiece(board.Data[taken]) != PAWN || GetPiece(board.Data[to]) != EMPTY || GetPiece(board.Data[from]) != PAWN {
		t.Fail()
	}
}
<|endoftext|>


<commit_before>package mp

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"github.com/ridewindx/mel"
	"github.com/ridewindx/mel/binding"
	"sort"
	"sync"
	"sync/atomic"
	"unsafe"
	"go.uber.org/zap"
	"github.com/jiudaoyun/wechat"
	"net/http"
	"fmt"
	"encoding/json"
	"strings"
)

type Server struct {
	*mel.Mel
	urlPrefix string

	appID string // App ID
	ID    string // Wechat ID

	tokenMutex sync.Mutex
	token      unsafe.Pointer

	aesKeyMutex sync.Mutex
	aesKey      unsafe.Pointer

	client *Client
	middlewares       []Handler
	messageHandlerMap map[string]Handler
	eventHandlerMap   map[string]Handler

	*zap.SugaredLogger
}

func (srv *Server) SetURLPrefix(urlPrefix string) {
	if !strings.HasPrefix(urlPrefix, "/") {
		urlPrefix = "/" + urlPrefix
	}
	urlPrefix = strings.TrimRight(urlPrefix, "/")
	srv.urlPrefix = urlPrefix
}

func (srv *Server) SetID(id string) {
	srv.ID = id
}

func (srv *Server) SetAppID(appID string) {
	srv.appID = appID
}

func (srv *Server) SetClient(client *Client) {
	srv.client = client
}

type Token struct {
	current string
	last    string
}

type AESKey struct {
	current string
	last    string
}

func (srv *Server) GetToken() (string, string) {
	p := (*Token)(atomic.LoadPointer(&srv.token))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetToken(token string) {
	if token == "" {
		return
	}

	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, _ := srv.GetToken()
	if token == current {
		return
	}

	t := Token{
		current: token,
		last:    current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) deleteLastToken() {
	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, last := srv.GetToken()
	if last == "" {
		return
	}

	t := Token{
		current: current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) GetAESKey() (string, string) {
	p := (*AESKey)(atomic.LoadPointer(&srv.aesKey))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetAESKey(base64AESKey string) {
	if len(base64AESKey) != 43 {
		return
	}
	aesKey, err := base64.StdEncoding.DecodeString(base64AESKey + "=")
	if err != nil {
		return
	}

	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, _ := srv.GetAESKey()
	if bytes.Equal(aesKey, []byte(current)) {
		return
	}

	k := AESKey{
		current: string(aesKey),
		last:    current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) deleteLastAESKey() {
	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, last := srv.GetAESKey()
	if last == "" {
		return
	}

	k := AESKey{
		current: current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) Use(middlewares ...Handler) {
	srv.middlewares = append(srv.middlewares, middlewares...)
	if len(srv.middlewares)+1 > int(abortIndex) {
		panic("too many middlewares")
	}
}

func (srv *Server) HandleMessage(msgType string, handler Handler) {
	srv.messageHandlerMap[msgType] = handler
}

func (srv *Server) HandleEvent(eventType string, handler Handler) {
	srv.eventHandlerMap[eventType] = handler
}

func NewServer(token, aesKey string) *Server {
	srv := &Server{
		Mel:               mel.New(),
		messageHandlerMap: make(map[string]Handler),
		eventHandlerMap:   make(map[string]Handler),
		SugaredLogger: wechat.Sugar,
	}

	srv.SetToken(token)
	srv.SetAESKey(aesKey)

	equal := func(a, b string) bool {
		return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
	}

	verifySignReturnToken := func(signature, timestamp, nonce string) string {
		currentToken, lastToken := srv.GetToken()
		token := currentToken

		isValid := func() bool {
			computedSignature := computeSign(token, timestamp, nonce)
			return equal(signature, computedSignature)
		}

		if isValid() {
			srv.deleteLastToken()
			return token
		}

		if lastToken != "" {
			token = lastToken
			if isValid() {
				return token
			}
		}

		return ""
	}

	verifySign := func(c *mel.Context) bool {
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		return verifySignReturnToken(signature, timestamp, nonce) != ""
	}

	type EncryptMsg struct {
		ToUserName string `xml:",cdata"`
		Encrypt    string `xml:",cdata"`
	}

	if srv.urlPrefix != "" {
		srv.Get("/", func(c *mel.Context) { // health check
			c.Status(200)
		})
	}

	srv.Get(srv.urlPrefix+"/", func(c *mel.Context) {
		if verifySign(c) {
			echostr := c.Query("echostr")
			c.Text(200, echostr)
		}
	})

	handleMessage := func(event *Event) interface{} {
		var handler Handler
		var ok bool
		if event.Type == MessageEvent {
			handler, ok = srv.eventHandlerMap[event.Event]
		} else {
			handler, ok = srv.messageHandlerMap[event.Type]
		}
		if !ok {
			return nil // no registered handler, just respond with empty string
		}

		ctx := &Context{
			Client: srv.client,
			index:    preStartIndex,
			handlers: append(srv.middlewares, handler),
			Event:    event,
		}

		ctx.Next()

		return ctx.response
	}

	srv.Post(srv.urlPrefix+"/", func(c *mel.Context) {
		encryptType := c.Query("encrypt_type")
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		switch encryptType {
		case "aes":
			token := verifySignReturnToken(signature, timestamp, nonce)
			if token == "" {
				return
			}

			msgSign := c.Query("msg_signature")

			var obj EncryptMsg
			err := c.BindWith(&obj, binding.XML)
			if err != nil {
				return
			}

			if srv.ID != "" && !equal(obj.ToUserName, srv.ID) {
				return
			}

			computedSign := computeSign(token, timestamp, nonce, obj.Encrypt)
			if !equal(computedSign, msgSign) {
				return
			}

			encryptedMsg, err := base64.StdEncoding.DecodeString(obj.Encrypt)
			if err != nil {
				return
			}

			current, last := srv.GetAESKey()
			aesKey := current
			random, msg, appId, err := decryptMsg(encryptedMsg, []byte(aesKey))
			if err != nil {
				if last == "" {
					return
				}
				aesKey = last
				random, msg, appId, err = decryptMsg(encryptedMsg, []byte(aesKey))
				if err != nil {
					return
				}
			} else {
				srv.deleteLastAESKey()
			}
			if srv.appID != "" && string(appId) != srv.appID {
				return
			}

			var event Event
			if err = xml.Unmarshal(msg, &event); err != nil {
				return
			}

			repBytes, err := xml.Marshal(handleMessage(&event))
			if err != nil {
				return
			}

			encryptedRepBytes := encryptMsg(random, repBytes, appId, []byte(aesKey))
			encryptedRepStr := base64.StdEncoding.EncodeToString(encryptedRepBytes)
			repSignature := computeSign(token, timestamp, nonce, encryptedRepStr)

			type EncryptRepMsg struct {
				Encrypt      string
				MsgSignature string
				TimeStamp    string
				Nonce        string
			}

			c.XML(200, &EncryptRepMsg{encryptedRepStr, repSignature, timestamp, nonce})

		case "", "raw":
			if !verifySign(c) {
				return
			}

			var event Event
			err := c.BindWith(&event, binding.XML)
			if err != nil {
				return
			}

			c.XML(200, handleMessage(&event))

		default:
			return
		}
	})

	handleAuthorize := func(c *mel.Context, url string, state string) {
		rep, err := srv.client.Client.Get(url)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
		defer rep.Body.Close()

		if rep.StatusCode != http.StatusOK {
			c.AbortWithError(http.StatusUnauthorized, fmt.Errorf("http.Status: %s", rep.Status))
			return
		}

		type Result struct {
			AccessToken string `json:"access_token"`
			ExpiresIn string `json:"expires_in"`
			RefreshToken string `json:"refresh_token"`
			OpenID string `json:"openid"`
			Scope string `json:"scope"`
			State string `json:"state,omitempty"`
		}

		type ResultWithErr struct {
			Result
			Err
		}

		var result ResultWithErr
		err = json.NewDecoder(rep.Body).Decode(&result)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}

		if result.Code() != OK {
			c.AbortWithError(http.StatusUnauthorized, &result)
			return
		}

		result.State = state
		c.JSON(http.StatusOK, &result.Result)
	}

	srv.Get(srv.urlPrefix+"/token", func(c *mel.Context) {
		code := c.Query("code")
		state := c.Query("state")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/access_token?appid=%s&secret=%s&code=%s&grant_type=authorization_code", srv.client.appId, srv.client.appSecret, code)

		handleAuthorize(c, url, state)
	})

	srv.Get(srv.urlPrefix+"/refresh-token", func(c *mel.Context) {
		refreshToken := c.Query("refresh_token")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/refresh_token?appid=%s&grant_type=refresh_token&refresh_token=%s", srv.client.appId, refreshToken)

		handleAuthorize(c, url, "")
	})

	srv.Get(srv.urlPrefix+"/signature", func(c *mel.Context) {
		timestamp := c.Query("timestamp")
		nonceStr := c.Query("nonceStr")
		url := c.Query("url")
		refresh := c.Query("refresh")

		var ticket string
		var err error
		if refresh != "" && (refresh == "true" || refresh == "True" || refresh == "1") {
			ticket, err = srv.client.RefreshTicket("")
		} else {
			ticket, err = srv.client.Ticket()
		}
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}

		strs := sort.StringSlice{
			"timestamp="+timestamp,
			"nonceStr="+nonceStr,
			"url="+url,
			"jsapi_ticket="+ticket,
		}
		strs.Sort()
		h := sha1.New()
		buf := bufio.NewWriterSize(h, 1024)
		for i, s := range strs {
			buf.WriteString(s)
			if i < len(strs)-1 {
				buf.WriteByte('&')
			}
		}
		buf.Flush()
		sign := hex.EncodeToString(h.Sum(nil))
		c.JSON(http.StatusOK, map[string]string{
			"signature": sign,
		})
	})

	return srv
}

func computeSign(elements ...string) string {
	strs := sort.StringSlice(elements)
	strs.Sort()

	h := sha1.New()

	buf := bufio.NewWriterSize(h, 1024)
	for _, s := range strs {
		buf.WriteString(s)
	}
	buf.Flush()

	return hex.EncodeToString(h.Sum(nil))
}
<commit_msg>Update source files<commit_after>package mp

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"github.com/ridewindx/mel"
	"github.com/ridewindx/mel/binding"
	"sort"
	"sync"
	"sync/atomic"
	"unsafe"
	"go.uber.org/zap"
	"github.com/jiudaoyun/wechat"
	"net/http"
	"fmt"
	"encoding/json"
	"strings"
)

type Server struct {
	*mel.Mel
	urlPrefix string

	appID string // App ID
	ID    string // Wechat ID

	tokenMutex sync.Mutex
	token      unsafe.Pointer

	aesKeyMutex sync.Mutex
	aesKey      unsafe.Pointer

	client *Client
	middlewares       []Handler
	messageHandlerMap map[string]Handler
	eventHandlerMap   map[string]Handler

	*zap.SugaredLogger
}

func (srv *Server) SetURLPrefix(urlPrefix string) {
	if !strings.HasPrefix(urlPrefix, "/") {
		urlPrefix = "/" + urlPrefix
	}
	urlPrefix = strings.TrimRight(urlPrefix, "/")
	srv.urlPrefix = urlPrefix
}

func (srv *Server) SetID(id string) {
	srv.ID = id
}

func (srv *Server) SetAppID(appID string) {
	srv.appID = appID
}

func (srv *Server) SetClient(client *Client) {
	srv.client = client
}

type Token struct {
	current string
	last    string
}

type AESKey struct {
	current string
	last    string
}

func (srv *Server) GetToken() (string, string) {
	p := (*Token)(atomic.LoadPointer(&srv.token))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetToken(token string) {
	if token == "" {
		return
	}

	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, _ := srv.GetToken()
	if token == current {
		return
	}

	t := Token{
		current: token,
		last:    current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) deleteLastToken() {
	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, last := srv.GetToken()
	if last == "" {
		return
	}

	t := Token{
		current: current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) GetAESKey() (string, string) {
	p := (*AESKey)(atomic.LoadPointer(&srv.aesKey))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetAESKey(base64AESKey string) {
	if len(base64AESKey) != 43 {
		return
	}
	aesKey, err := base64.StdEncoding.DecodeString(base64AESKey + "=")
	if err != nil {
		return
	}

	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, _ := srv.GetAESKey()
	if bytes.Equal(aesKey, []byte(current)) {
		return
	}

	k := AESKey{
		current: string(aesKey),
		last:    current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) deleteLastAESKey() {
	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, last := srv.GetAESKey()
	if last == "" {
		return
	}

	k := AESKey{
		current: current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) Use(middlewares ...Handler) {
	srv.middlewares = append(srv.middlewares, middlewares...)
	if len(srv.middlewares)+1 > int(abortIndex) {
		panic("too many middlewares")
	}
}

func (srv *Server) HandleMessage(msgType string, handler Handler) {
	srv.messageHandlerMap[msgType] = handler
}

func (srv *Server) HandleEvent(eventType string, handler Handler) {
	srv.eventHandlerMap[eventType] = handler
}

func NewServer(token, aesKey string) *Server {
	srv := &Server{
		Mel:               mel.New(),
		messageHandlerMap: make(map[string]Handler),
		eventHandlerMap:   make(map[string]Handler),
		SugaredLogger: wechat.Sugar,
	}

	srv.SetToken(token)
	srv.SetAESKey(aesKey)

	equal := func(a, b string) bool {
		return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
	}

	verifySignReturnToken := func(signature, timestamp, nonce string) string {
		currentToken, lastToken := srv.GetToken()
		token := currentToken

		isValid := func() bool {
			computedSignature := computeSign(token, timestamp, nonce)
			return equal(signature, computedSignature)
		}

		if isValid() {
			srv.deleteLastToken()
			return token
		}

		if lastToken != "" {
			token = lastToken
			if isValid() {
				return token
			}
		}

		return ""
	}

	verifySign := func(c *mel.Context) bool {
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		return verifySignReturnToken(signature, timestamp, nonce) != ""
	}

	type EncryptMsg struct {
		ToUserName string `xml:",cdata"`
		Encrypt    string `xml:",cdata"`
	}

	srv.Head("/", func(c *mel.Context) { // health check
		c.Status(200)
	})

	srv.Get(srv.urlPrefix+"/", func(c *mel.Context) {
		if verifySign(c) {
			echostr := c.Query("echostr")
			c.Text(200, echostr)
		}
	})

	handleMessage := func(event *Event) interface{} {
		var handler Handler
		var ok bool
		if event.Type == MessageEvent {
			handler, ok = srv.eventHandlerMap[event.Event]
		} else {
			handler, ok = srv.messageHandlerMap[event.Type]
		}
		if !ok {
			return nil // no registered handler, just respond with empty string
		}

		ctx := &Context{
			Client: srv.client,
			index:    preStartIndex,
			handlers: append(srv.middlewares, handler),
			Event:    event,
		}

		ctx.Next()

		return ctx.response
	}

	srv.Post(srv.urlPrefix+"/", func(c *mel.Context) {
		encryptType := c.Query("encrypt_type")
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		switch encryptType {
		case "aes":
			token := verifySignReturnToken(signature, timestamp, nonce)
			if token == "" {
				return
			}

			msgSign := c.Query("msg_signature")

			var obj EncryptMsg
			err := c.BindWith(&obj, binding.XML)
			if err != nil {
				return
			}

			if srv.ID != "" && !equal(obj.ToUserName, srv.ID) {
				return
			}

			computedSign := computeSign(token, timestamp, nonce, obj.Encrypt)
			if !equal(computedSign, msgSign) {
				return
			}

			encryptedMsg, err := base64.StdEncoding.DecodeString(obj.Encrypt)
			if err != nil {
				return
			}

			current, last := srv.GetAESKey()
			aesKey := current
			random, msg, appId, err := decryptMsg(encryptedMsg, []byte(aesKey))
			if err != nil {
				if last == "" {
					return
				}
				aesKey = last
				random, msg, appId, err = decryptMsg(encryptedMsg, []byte(aesKey))
				if err != nil {
					return
				}
			} else {
				srv.deleteLastAESKey()
			}
			if srv.appID != "" && string(appId) != srv.appID {
				return
			}

			var event Event
			if err = xml.Unmarshal(msg, &event); err != nil {
				return
			}

			repBytes, err := xml.Marshal(handleMessage(&event))
			if err != nil {
				return
			}

			encryptedRepBytes := encryptMsg(random, repBytes, appId, []byte(aesKey))
			encryptedRepStr := base64.StdEncoding.EncodeToString(encryptedRepBytes)
			repSignature := computeSign(token, timestamp, nonce, encryptedRepStr)

			type EncryptRepMsg struct {
				Encrypt      string
				MsgSignature string
				TimeStamp    string
				Nonce        string
			}

			c.XML(200, &EncryptRepMsg{encryptedRepStr, repSignature, timestamp, nonce})

		case "", "raw":
			if !verifySign(c) {
				return
			}

			var event Event
			err := c.BindWith(&event, binding.XML)
			if err != nil {
				return
			}

			c.XML(200, handleMessage(&event))

		default:
			return
		}
	})

	handleAuthorize := func(c *mel.Context, url string, state string) {
		rep, err := srv.client.Client.Get(url)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
		defer rep.Body.Close()

		if rep.StatusCode != http.StatusOK {
			c.AbortWithError(http.StatusUnauthorized, fmt.Errorf("http.Status: %s", rep.Status))
			return
		}

		type Result struct {
			AccessToken string `json:"access_token"`
			ExpiresIn string `json:"expires_in"`
			RefreshToken string `json:"refresh_token"`
			OpenID string `json:"openid"`
			Scope string `json:"scope"`
			State string `json:"state,omitempty"`
		}

		type ResultWithErr struct {
			Result
			Err
		}

		var result ResultWithErr
		err = json.NewDecoder(rep.Body).Decode(&result)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}

		if result.Code() != OK {
			c.AbortWithError(http.StatusUnauthorized, &result)
			return
		}

		result.State = state
		c.JSON(http.StatusOK, &result.Result)
	}

	srv.Get(srv.urlPrefix+"/token", func(c *mel.Context) {
		code := c.Query("code")
		state := c.Query("state")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/access_token?appid=%s&secret=%s&code=%s&grant_type=authorization_code", srv.client.appId, srv.client.appSecret, code)

		handleAuthorize(c, url, state)
	})

	srv.Get(srv.urlPrefix+"/refresh-token", func(c *mel.Context) {
		refreshToken := c.Query("refresh_token")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/refresh_token?appid=%s&grant_type=refresh_token&refresh_token=%s", srv.client.appId, refreshToken)

		handleAuthorize(c, url, "")
	})

	srv.Get(srv.urlPrefix+"/signature", func(c *mel.Context) {
		timestamp := c.Query("timestamp")
		nonceStr := c.Query("nonceStr")
		url := c.Query("url")
		refresh := c.Query("refresh")

		var ticket string
		var err error
		if refresh != "" && (refresh == "true" || refresh == "True" || refresh == "1") {
			ticket, err = srv.client.RefreshTicket("")
		} else {
			ticket, err = srv.client.Ticket()
		}
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}

		strs := sort.StringSlice{
			"timestamp="+timestamp,
			"nonceStr="+nonceStr,
			"url="+url,
			"jsapi_ticket="+ticket,
		}
		strs.Sort()
		h := sha1.New()
		buf := bufio.NewWriterSize(h, 1024)
		for i, s := range strs {
			buf.WriteString(s)
			if i < len(strs)-1 {
				buf.WriteByte('&')
			}
		}
		buf.Flush()
		sign := hex.EncodeToString(h.Sum(nil))
		c.JSON(http.StatusOK, map[string]string{
			"signature": sign,
		})
	})

	return srv
}

func computeSign(elements ...string) string {
	strs := sort.StringSlice(elements)
	strs.Sort()

	h := sha1.New()

	buf := bufio.NewWriterSize(h, 1024)
	for _, s := range strs {
		buf.WriteString(s)
	}
	buf.Flush()

	return hex.EncodeToString(h.Sum(nil))
}
<|endoftext|>