<commit_before>package goredis

import (
	"errors"
	"io"
	"net"
	"strconv"
)

// Instruct Redis to start an Append Only File rewrite process.
// The rewrite will create a small optimized version of the current Append Only File.
func (r *Redis) BgRewriteAof() error {
	_, err := r.ExecuteCommand("BGREWRITEAOF")
	return err
}

// Save the DB in background.
// The OK code is immediately returned.
// Redis forks, the parent continues to serve the clients, the child saves the DB on disk then exits.
// A client my be able to check if the operation succeeded using the LASTSAVE command.
func (r *Redis) BgSave() error {
	_, err := r.ExecuteCommand("BGSAVE")
	return err
}

// The CLIENT KILL command closes a given client connection identified by ip:port.
// Due to the single-treaded nature of Redis, it is not possible to kill a client connection while it is executing a command.
// However, the client will notice the connection has been closed only when the next command is sent (and results in network error).
// Status code reply: OK if the connection exists and has been closed
func (r *Redis) ClientKill(ip string, port int) error {
	rp, err := r.ExecuteCommand("CLIENT", "KILL", net.JoinHostPort(ip, strconv.Itoa(port)))
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// The CLIENT LIST command returns information and statistics about the client connections server in a mostly human readable format.
// Bulk reply: a unique string, formatted as follows:
// One client connection per line (separated by LF)
// Each line is composed of a succession of property=value fields separated by a space character.
func (r *Redis) ClientList() (string, error) {
	rp, err := r.ExecuteCommand("CLIENT", "LIST")
	if err != nil {
		return "", err
	}
	return rp.StringValue()
}

// The CLIENT GETNAME returns the name of the current connection as set by CLIENT SETNAME.
// Since every new connection starts without an associated name,
// if no name was assigned a null bulk reply is returned.
func (r *Redis) ClientGetName() ([]byte, error) {
	rp, err := r.ExecuteCommand("CLIENT", "GETNAME")
	if err != nil {
		return nil, err
	}
	return rp.BytesValue()
}

// The CLIENT SETNAME command assigns a name to the current connection.
func (r *Redis) ClientSetName(name string) error {
	rp, err := r.ExecuteCommand("CLIENT", "SETNAME", name)
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// The CONFIG GET command is used to read the configuration parameters of a running Redis server.
// Not all the configuration parameters are supported in Redis 2.4,
// while Redis 2.6 can read the whole configuration of a server using this command.
// CONFIG GET takes a single argument, which is a glob-style pattern.
func (r *Redis) ConfigGet(parameter string) (map[string]string, error) {
	rp, err := r.ExecuteCommand("CONFIG", "GET", parameter)
	if err != nil {
		return nil, err
	}
	return rp.HashValue()
}

// The CONFIG REWRITE command rewrites the redis.conf file the server was started with,
// applying the minimal changes needed to make it reflecting the configuration currently used by the server,
// that may be different compared to the original one because of the use of the CONFIG SET command.
// Available since 2.8.0.
func (r *Redis) ConfigRewrite() error {
	rp, err := r.ExecuteCommand("CONFIG", "REWRITE")
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// The CONFIG SET command is used in order to reconfigure the server at run time without the need to restart Redis.
// You can change both trivial parameters or switch from one to another persistence option using this command.
func (r *Redis) ConfigSet(parameter, value string) error {
	rp, err := r.ExecuteCommand("CONFIG", "SET")
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// Resets the statistics reported by Redis using the INFO command.
// These are the counters that are reset:
// Keyspace hits
// Keyspace misses
// Number of commands processed
// Number of connections received
// Number of expired keys
// Number of rejected connections
// Latest fork(2) time
// The aof_delayed_fsync counter
func (r *Redis) ConfigResetStat() error {
	_, err := r.ExecuteCommand("CONFIG", "RESETSTAT")
	return err
}

// Return the number of keys in the currently-selected database.
func (r *Redis) DBSize() (int64, error) {
	rp, err := r.ExecuteCommand("DBSIZE")
	if err != nil {
		return 0, err
	}
	return rp.IntegerValue()
}

// DEBUG OBJECT key
// DEBUG OBJECT is a debugging command that should not be used by clients.
func (r *Redis) DebugObject(key string) (string, error) {
	rp, err := r.ExecuteCommand("DEBUG", "OBJECT", key)
	if err != nil {
		return "", err
	}
	return rp.StatusValue()
}

// Delete all the keys of all the existing databases,
// not just the currently selected one.
// This command never fails.
func (r *Redis) FlushAll() error {
	_, err := r.ExecuteCommand("FLUSHALL")
	return err
}

// Delete all the keys of the currently selected DB.
// This command never fails.
func (r *Redis) FlushDB() error {
	_, err := r.ExecuteCommand("FLUSHDB")
	return err
}

// The INFO command returns information and statistics about the server
// in a format that is simple to parse by computers and easy to read by humans.
// format document at http://redis.io/commands/info
func (r *Redis) Info(section string) (string, error) {
	args := packArgs("INFO", section)
	rp, err := r.ExecuteCommand(args...)
	if err != nil {
		return "", err
	}
	return rp.StringValue()
}

// Return the UNIX TIME of the last DB save executed with success.
// A client may check if a BGSAVE command succeeded reading the LASTSAVE value,
// then issuing a BGSAVE command and checking at regular intervals every N seconds if LASTSAVE changed.
// Integer reply: an UNIX time stamp.
func (r *Redis) LastSave() (int64, error) {
	rp, err := r.ExecuteCommand("LASTSAVE")
	if err != nil {
		return 0, err
	}
	return rp.IntegerValue()
}

// MONITOR is a debugging command that streams back every command processed by the Redis server.
type MonitorCommand struct {
	redis *Redis
	conn  *Connection
}

func (r *Redis) Monitor() (*MonitorCommand, error) {
	c, err := r.openConnection()
	if err != nil {
		return nil, err
	}
	if err := c.SendCommand("MONITOR"); err != nil {
		return nil, err
	}
	rp, err := c.RecvReply()
	if err != nil {
		return nil, err
	}
	if err := rp.OKValue(); err != nil {
		return nil, err
	}
	return &MonitorCommand{r, c}, nil
}

func (m *MonitorCommand) Receive() (string, error) {
	rp, err := m.conn.RecvReply()
	if err != nil {
		return "", err
	}
	return rp.StatusValue()
}

func (m *MonitorCommand) Close() error {
	return m.conn.SendCommand("QUIT")
}

// The SAVE commands performs a synchronous save of the dataset producing a point in time snapshot of all the data inside the Redis instance,
// in the form of an RDB file.
//
// You almost never want to call SAVE in production environments where it will block all the other clients. Instead usually BGSAVE is used.
func (r *Redis) Save() error {
	rp, err := r.ExecuteCommand("SAVE")
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// The command behavior is the following:
// Stop all the clients.
// Perform a blocking SAVE if at least one save point is configured.
// Flush the Append Only File if AOF is enabled.
// Quit the server.
func (r *Redis) Shutdown(save, no_save bool) error {
	args := packArgs("SHUTDOWN")
	if save {
		args = append(args, "SAVE")
	} else if no_save {
		args = append(args, "NOSAVE")
	}
	rp, err := r.ExecuteCommand(args...)
	if err == io.EOF {
		return nil
	}
	if err != nil {
		return err
	}
	return errors.New(rp.Status)
}

// The SLAVEOF command can change the replication settings of a slave on the fly.
// If a Redis server is already acting as slave, the command SLAVEOF NO ONE will turn off the replication,
// turning the Redis server into a MASTER.
// In the proper form SLAVEOF hostname port will make the server a slave of another server listening at the specified hostname and port.
//
// If a server is already a slave of some master,
// SLAVEOF hostname port will stop the replication against the old server and start the synchronization against the new one, discarding the old dataset.
// The form SLAVEOF NO ONE will stop replication, turning the server into a MASTER, but will not discard the replication.
// So, if the old master stops working, it is possible to turn the slave into a master and set the application to use this new master in read/write.
// Later when the other Redis server is fixed, it can be reconfigured to work as a slave.
func (r *Redis) SlaveOf(host, port string) error {
	rp, err := r.ExecuteCommand("SLAVEOF", host, port)
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// SLOWLOG subcommand [argument]
// This command is used in order to read and reset the Redis slow queries log.
type SlowLog struct {
	ID           int64
	Timestamp    int64
	Microseconds int64
	Command      []string
}

func (r *Redis) SlowLogGet(n int) ([]*SlowLog, error) {
	rp, err := r.ExecuteCommand("SLOWLOG", "GET", n)
	if err != nil {
		return nil, err
	}
	if rp.Type == ErrorReply {
		return nil, errors.New(rp.Error)
	}
	if rp.Type != MultiReply {
		return nil, errors.New("slowlog get protocol error")
	}
	var slow []*SlowLog
	for _, subrp := range rp.Multi {
		if subrp.Multi == nil || len(subrp.Multi) != 4 {
			return nil, errors.New("slowlog get protocol error")
		}
		id, err := subrp.Multi[0].IntegerValue()
		if err != nil {
			return nil, err
		}
		timestamp, err := subrp.Multi[1].IntegerValue()
		if err != nil {
			return nil, err
		}
		microseconds, err := subrp.Multi[2].IntegerValue()
		if err != nil {
			return nil, err
		}
		command, err := subrp.Multi[3].ListValue()
		if err != nil {
			return nil, err
		}
		slow = append(slow, &SlowLog{id, timestamp, microseconds, command})
	}
	return slow, nil
}

// Obtaining the current length of the slow log
func (r *Redis) SlowLogLen() (int64, error) {
	rp, err := r.ExecuteCommand("SLOWLOG", "LEN")
	if err != nil {
		return 0, err
	}
	return rp.IntegerValue()
}

// Resetting the slow log. Once deleted the information is lost forever.
func (r *Redis) SlowLogReset() error {
	rp, err := r.ExecuteCommand("SLOWLOG", "RESET")
	if err != nil {
		return err
	}
	return rp.OKValue()
}

/*
SYNC Internal command used for replication
*/

// A multi bulk reply containing two elements:
// unix time in seconds.
// microseconds.
func (r *Redis) Time() ([]string, error) {
	rp, err := r.ExecuteCommand("TIME")
	if err != nil {
		return nil, err
	}
	return rp.ListValue()
}
<commit_msg>server: sync command will never be imeplemented<commit_after>package goredis

import (
	"errors"
	"io"
	"net"
	"strconv"
)

// Instruct Redis to start an Append Only File rewrite process.
// The rewrite will create a small optimized version of the current Append Only File.
func (r *Redis) BgRewriteAof() error {
	_, err := r.ExecuteCommand("BGREWRITEAOF")
	return err
}

// Save the DB in background.
// The OK code is immediately returned.
// Redis forks, the parent continues to serve the clients, the child saves the DB on disk then exits.
// A client my be able to check if the operation succeeded using the LASTSAVE command.
func (r *Redis) BgSave() error {
	_, err := r.ExecuteCommand("BGSAVE")
	return err
}

// The CLIENT KILL command closes a given client connection identified by ip:port.
// Due to the single-treaded nature of Redis, it is not possible to kill a client connection while it is executing a command.
// However, the client will notice the connection has been closed only when the next command is sent (and results in network error).
// Status code reply: OK if the connection exists and has been closed
func (r *Redis) ClientKill(ip string, port int) error {
	rp, err := r.ExecuteCommand("CLIENT", "KILL", net.JoinHostPort(ip, strconv.Itoa(port)))
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// The CLIENT LIST command returns information and statistics about the client connections server in a mostly human readable format.
// Bulk reply: a unique string, formatted as follows:
// One client connection per line (separated by LF)
// Each line is composed of a succession of property=value fields separated by a space character.
func (r *Redis) ClientList() (string, error) {
	rp, err := r.ExecuteCommand("CLIENT", "LIST")
	if err != nil {
		return "", err
	}
	return rp.StringValue()
}

// The CLIENT GETNAME returns the name of the current connection as set by CLIENT SETNAME.
// Since every new connection starts without an associated name,
// if no name was assigned a null bulk reply is returned.
func (r *Redis) ClientGetName() ([]byte, error) {
	rp, err := r.ExecuteCommand("CLIENT", "GETNAME")
	if err != nil {
		return nil, err
	}
	return rp.BytesValue()
}

// The CLIENT SETNAME command assigns a name to the current connection.
func (r *Redis) ClientSetName(name string) error {
	rp, err := r.ExecuteCommand("CLIENT", "SETNAME", name)
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// The CONFIG GET command is used to read the configuration parameters of a running Redis server.
// Not all the configuration parameters are supported in Redis 2.4,
// while Redis 2.6 can read the whole configuration of a server using this command.
// CONFIG GET takes a single argument, which is a glob-style pattern.
func (r *Redis) ConfigGet(parameter string) (map[string]string, error) {
	rp, err := r.ExecuteCommand("CONFIG", "GET", parameter)
	if err != nil {
		return nil, err
	}
	return rp.HashValue()
}

// The CONFIG REWRITE command rewrites the redis.conf file the server was started with,
// applying the minimal changes needed to make it reflecting the configuration currently used by the server,
// that may be different compared to the original one because of the use of the CONFIG SET command.
// Available since 2.8.0.
func (r *Redis) ConfigRewrite() error {
	rp, err := r.ExecuteCommand("CONFIG", "REWRITE")
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// The CONFIG SET command is used in order to reconfigure the server at run time without the need to restart Redis.
// You can change both trivial parameters or switch from one to another persistence option using this command.
func (r *Redis) ConfigSet(parameter, value string) error {
	rp, err := r.ExecuteCommand("CONFIG", "SET")
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// Resets the statistics reported by Redis using the INFO command.
// These are the counters that are reset:
// Keyspace hits
// Keyspace misses
// Number of commands processed
// Number of connections received
// Number of expired keys
// Number of rejected connections
// Latest fork(2) time
// The aof_delayed_fsync counter
func (r *Redis) ConfigResetStat() error {
	_, err := r.ExecuteCommand("CONFIG", "RESETSTAT")
	return err
}

// Return the number of keys in the currently-selected database.
func (r *Redis) DBSize() (int64, error) {
	rp, err := r.ExecuteCommand("DBSIZE")
	if err != nil {
		return 0, err
	}
	return rp.IntegerValue()
}

// DEBUG OBJECT key
// DEBUG OBJECT is a debugging command that should not be used by clients.
func (r *Redis) DebugObject(key string) (string, error) {
	rp, err := r.ExecuteCommand("DEBUG", "OBJECT", key)
	if err != nil {
		return "", err
	}
	return rp.StatusValue()
}

// Delete all the keys of all the existing databases,
// not just the currently selected one.
// This command never fails.
func (r *Redis) FlushAll() error {
	_, err := r.ExecuteCommand("FLUSHALL")
	return err
}

// Delete all the keys of the currently selected DB.
// This command never fails.
func (r *Redis) FlushDB() error {
	_, err := r.ExecuteCommand("FLUSHDB")
	return err
}

// The INFO command returns information and statistics about the server
// in a format that is simple to parse by computers and easy to read by humans.
// format document at http://redis.io/commands/info
func (r *Redis) Info(section string) (string, error) {
	args := packArgs("INFO", section)
	rp, err := r.ExecuteCommand(args...)
	if err != nil {
		return "", err
	}
	return rp.StringValue()
}

// Return the UNIX TIME of the last DB save executed with success.
// A client may check if a BGSAVE command succeeded reading the LASTSAVE value,
// then issuing a BGSAVE command and checking at regular intervals every N seconds if LASTSAVE changed.
// Integer reply: an UNIX time stamp.
func (r *Redis) LastSave() (int64, error) {
	rp, err := r.ExecuteCommand("LASTSAVE")
	if err != nil {
		return 0, err
	}
	return rp.IntegerValue()
}

// MONITOR is a debugging command that streams back every command processed by the Redis server.
type MonitorCommand struct {
	redis *Redis
	conn  *Connection
}

func (r *Redis) Monitor() (*MonitorCommand, error) {
	c, err := r.openConnection()
	if err != nil {
		return nil, err
	}
	if err := c.SendCommand("MONITOR"); err != nil {
		return nil, err
	}
	rp, err := c.RecvReply()
	if err != nil {
		return nil, err
	}
	if err := rp.OKValue(); err != nil {
		return nil, err
	}
	return &MonitorCommand{r, c}, nil
}

func (m *MonitorCommand) Receive() (string, error) {
	rp, err := m.conn.RecvReply()
	if err != nil {
		return "", err
	}
	return rp.StatusValue()
}

func (m *MonitorCommand) Close() error {
	return m.conn.SendCommand("QUIT")
}

// The SAVE commands performs a synchronous save of the dataset producing a point in time snapshot of all the data inside the Redis instance,
// in the form of an RDB file.
//
// You almost never want to call SAVE in production environments where it will block all the other clients. Instead usually BGSAVE is used.
func (r *Redis) Save() error {
	rp, err := r.ExecuteCommand("SAVE")
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// The command behavior is the following:
// Stop all the clients.
// Perform a blocking SAVE if at least one save point is configured.
// Flush the Append Only File if AOF is enabled.
// Quit the server.
func (r *Redis) Shutdown(save, no_save bool) error {
	args := packArgs("SHUTDOWN")
	if save {
		args = append(args, "SAVE")
	} else if no_save {
		args = append(args, "NOSAVE")
	}
	rp, err := r.ExecuteCommand(args...)
	if err == io.EOF {
		return nil
	}
	if err != nil {
		return err
	}
	return errors.New(rp.Status)
}

// The SLAVEOF command can change the replication settings of a slave on the fly.
// If a Redis server is already acting as slave, the command SLAVEOF NO ONE will turn off the replication,
// turning the Redis server into a MASTER.
// In the proper form SLAVEOF hostname port will make the server a slave of another server listening at the specified hostname and port.
//
// If a server is already a slave of some master,
// SLAVEOF hostname port will stop the replication against the old server and start the synchronization against the new one, discarding the old dataset.
// The form SLAVEOF NO ONE will stop replication, turning the server into a MASTER, but will not discard the replication.
// So, if the old master stops working, it is possible to turn the slave into a master and set the application to use this new master in read/write.
// Later when the other Redis server is fixed, it can be reconfigured to work as a slave.
func (r *Redis) SlaveOf(host, port string) error {
	rp, err := r.ExecuteCommand("SLAVEOF", host, port)
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// SLOWLOG subcommand [argument]
// This command is used in order to read and reset the Redis slow queries log.
type SlowLog struct {
	ID           int64
	Timestamp    int64
	Microseconds int64
	Command      []string
}

func (r *Redis) SlowLogGet(n int) ([]*SlowLog, error) {
	rp, err := r.ExecuteCommand("SLOWLOG", "GET", n)
	if err != nil {
		return nil, err
	}
	if rp.Type == ErrorReply {
		return nil, errors.New(rp.Error)
	}
	if rp.Type != MultiReply {
		return nil, errors.New("slowlog get protocol error")
	}
	var slow []*SlowLog
	for _, subrp := range rp.Multi {
		if subrp.Multi == nil || len(subrp.Multi) != 4 {
			return nil, errors.New("slowlog get protocol error")
		}
		id, err := subrp.Multi[0].IntegerValue()
		if err != nil {
			return nil, err
		}
		timestamp, err := subrp.Multi[1].IntegerValue()
		if err != nil {
			return nil, err
		}
		microseconds, err := subrp.Multi[2].IntegerValue()
		if err != nil {
			return nil, err
		}
		command, err := subrp.Multi[3].ListValue()
		if err != nil {
			return nil, err
		}
		slow = append(slow, &SlowLog{id, timestamp, microseconds, command})
	}
	return slow, nil
}

// Obtaining the current length of the slow log
func (r *Redis) SlowLogLen() (int64, error) {
	rp, err := r.ExecuteCommand("SLOWLOG", "LEN")
	if err != nil {
		return 0, err
	}
	return rp.IntegerValue()
}

// Resetting the slow log. Once deleted the information is lost forever.
func (r *Redis) SlowLogReset() error {
	rp, err := r.ExecuteCommand("SLOWLOG", "RESET")
	if err != nil {
		return err
	}
	return rp.OKValue()
}

// A multi bulk reply containing two elements:
// unix time in seconds.
// microseconds.
func (r *Redis) Time() ([]string, error) {
	rp, err := r.ExecuteCommand("TIME")
	if err != nil {
		return nil, err
	}
	return rp.ListValue()
}
<|endoftext|>


<commit_before>package proxy

import (
	"net/http"
)

func NewHandler(t *http.Transport, addrs []string) (http.Handler, error) {
	d, err := newDirector("http", addrs)
	if err != nil {
		return nil, err
	}

	rp := reverseProxy{
		director:  d,
		transport: t,
	}

	return &rp, nil
}

func readonlyHandlerFunc(next http.Handler) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		if req.Method != "GET" {
			w.WriteHeader(http.StatusNotImplemented)
			return
		}

		next.ServeHTTP(w, req)
	}
}

func NewReadonlyHandler(hdlr http.Handler) http.Handler {
	readonly := readonlyHandlerFunc(hdlr)
	return http.HandlerFunc(readonly)
}
<commit_msg>proxy: determine scheme based on TLSClientConfig<commit_after>package proxy

import (
	"net/http"
)

func NewHandler(t *http.Transport, addrs []string) (http.Handler, error) {
	scheme := "http"
	if t.TLSClientConfig != nil {
		scheme = "https"
	}

	d, err := newDirector(scheme, addrs)
	if err != nil {
		return nil, err
	}

	rp := reverseProxy{
		director:  d,
		transport: t,
	}

	return &rp, nil
}

func readonlyHandlerFunc(next http.Handler) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		if req.Method != "GET" {
			w.WriteHeader(http.StatusNotImplemented)
			return
		}

		next.ServeHTTP(w, req)
	}
}

func NewReadonlyHandler(hdlr http.Handler) http.Handler {
	readonly := readonlyHandlerFunc(hdlr)
	return http.HandlerFunc(readonly)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"strconv"
	"text/template"
)

var serverTmpl = new(template.Template)

func init() {
	serverTmpl.Funcs(template.FuncMap{
		"average": func() RollingAverage {
			return <-average
		},

		"format": func(num int64, base int) string {
			return strconv.FormatInt(num, base)
		},
	})

	template.Must(serverTmpl.New("main").Parse(`

		<title>{{.Title}} :: Main</title>

		{{with $avg := average -}}

Current average: {{$avg.Cur}}

Calculated from {{$avg.Num}} logouts.

		{{- end}}

`))
}

func logHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		log.Printf("Got %q request for %q", req.Method, req.URL)

		h.ServeHTTP(rw, req)
	})
}

func tmplHandler(t string) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		err := serverTmpl.ExecuteTemplate(rw, t, map[string]interface{}{
			"Req":   req,
			"Title": "PS2 Average Login Times",
		})
		if err != nil {
			log.Printf("Failed to execute %q: %v", t, err)
		}
	})
}

func serveAverage(rw http.ResponseWriter, req *http.Request) {
	e := json.NewEncoder(rw)
	err := e.Encode(<-average)
	if err != nil {
		log.Printf("Failed to write average: %v", err)
	}
}

func serveJS(rw http.ResponseWriter, req *http.Request) {
	_, err := io.WriteString(rw, `$(document).ready(function() {
	avg = $('#average');
	num = $('#num');
	function getAverage() {
		$.getJSON('/average', function(data) {
			avg.html(data.cur);
			num.html(data.num);
		});
	};

	setInterval(getAverage, 5000);
});`)
	if err != nil {
		log.Printf("Failed to write JS: %v", err)
	}
}

func server() {
	const (
		ServerAddr = ":8080"
	)

	http.Handle("/average", logHandler(http.HandlerFunc(serveAverage)))
	http.Handle("/ps2avglogin.js", logHandler(http.HandlerFunc(serveJS)))
	http.Handle("/", logHandler(tmplHandler("main")))

	log.Printf("Starting server at %q...", ServerAddr)
	err := http.ListenAndServe(ServerAddr, nil)
	if err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
}
<commit_msg>Fix incorrect close tag.<commit_after>package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"strconv"
	"text/template"
)

var serverTmpl = new(template.Template)

func init() {
	serverTmpl.Funcs(template.FuncMap{
		"average": func() RollingAverage {
			return <-average
		},

		"format": func(num int64, base int) string {
			return strconv.FormatInt(num, base)
		},
	})

	template.Must(serverTmpl.New("main").Parse(`

		<title>{{.Title}} :: Main</title>

		{{with $avg := average -}}

Current average: {{$avg.Cur}}

Calculated from {{$avg.Num}} logouts.

		{{- end}}

`))
}

func logHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		log.Printf("Got %q request for %q", req.Method, req.URL)

		h.ServeHTTP(rw, req)
	})
}

func tmplHandler(t string) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		err := serverTmpl.ExecuteTemplate(rw, t, map[string]interface{}{
			"Req":   req,
			"Title": "PS2 Average Login Times",
		})
		if err != nil {
			log.Printf("Failed to execute %q: %v", t, err)
		}
	})
}

func serveAverage(rw http.ResponseWriter, req *http.Request) {
	e := json.NewEncoder(rw)
	err := e.Encode(<-average)
	if err != nil {
		log.Printf("Failed to write average: %v", err)
	}
}

func serveJS(rw http.ResponseWriter, req *http.Request) {
	_, err := io.WriteString(rw, `$(document).ready(function() {
	avg = $('#average');
	num = $('#num');
	function getAverage() {
		$.getJSON('/average', function(data) {
			avg.html(data.cur);
			num.html(data.num);
		});
	};

	setInterval(getAverage, 5000);
});`)
	if err != nil {
		log.Printf("Failed to write JS: %v", err)
	}
}

func server() {
	const (
		ServerAddr = ":8080"
	)

	http.Handle("/average", logHandler(http.HandlerFunc(serveAverage)))
	http.Handle("/ps2avglogin.js", logHandler(http.HandlerFunc(serveJS)))
	http.Handle("/", logHandler(tmplHandler("main")))

	log.Printf("Starting server at %q...", ServerAddr)
	err := http.ListenAndServe(ServerAddr, nil)
	if err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
}
<|endoftext|>


<commit_before>// Package proxy is a cli proxy
package proxy

import (
	"os"
	"strings"
	"time"

	"github.com/go-acme/lego/v3/providers/dns/cloudflare"
	"github.com/micro/cli/v2"
	"github.com/micro/go-micro/v2"
	"github.com/micro/go-micro/v2/api/server/acme"
	"github.com/micro/go-micro/v2/api/server/acme/autocert"
	"github.com/micro/go-micro/v2/api/server/acme/certmagic"
	bmem "github.com/micro/go-micro/v2/broker/memory"
	"github.com/micro/go-micro/v2/client"
	mucli "github.com/micro/go-micro/v2/client"
	log "github.com/micro/go-micro/v2/logger"
	"github.com/micro/go-micro/v2/proxy"
	"github.com/micro/go-micro/v2/proxy/grpc"
	"github.com/micro/go-micro/v2/proxy/http"
	"github.com/micro/go-micro/v2/proxy/mucp"
	"github.com/micro/go-micro/v2/registry"
	rmem "github.com/micro/go-micro/v2/registry/memory"
	"github.com/micro/go-micro/v2/router"
	rs "github.com/micro/go-micro/v2/router/service"
	"github.com/micro/go-micro/v2/server"
	sgrpc "github.com/micro/go-micro/v2/server/grpc"
	cfstore "github.com/micro/go-micro/v2/store/cloudflare"
	"github.com/micro/go-micro/v2/sync/lock/memory"
	"github.com/micro/go-micro/v2/util/mux"
	"github.com/micro/micro/v2/internal/helper"
)

var (
	// Name of the proxy
	Name = "go.micro.proxy"
	// The address of the proxy
	Address = ":8081"
	// the proxy protocol
	Protocol = "grpc"
	// The endpoint host to route to
	Endpoint string
	// ACME (Cert management)
	ACMEProvider          = "autocert"
	ACMEChallengeProvider = "cloudflare"
	ACMECA                = acme.LetsEncryptProductionCA
)

func run(ctx *cli.Context, srvOpts ...micro.Option) {
	log.Init(log.WithFields(map[string]interface{}{"service": "proxy"}))

	// because MICRO_PROXY_ADDRESS is used internally by the go-micro/client
	// we need to unset it so we don't end up calling ourselves infinitely
	os.Unsetenv("MICRO_PROXY_ADDRESS")

	if len(ctx.String("server_name")) > 0 {
		Name = ctx.String("server_name")
	}
	if len(ctx.String("address")) > 0 {
		Address = ctx.String("address")
	}
	if len(ctx.String("endpoint")) > 0 {
		Endpoint = ctx.String("endpoint")
	}
	if len(ctx.String("protocol")) > 0 {
		Protocol = ctx.String("protocol")
	}
	if len(ctx.String("acme_provider")) > 0 {
		ACMEProvider = ctx.String("acme_provider")
	}

	// Init plugins
	for _, p := range Plugins() {
		p.Init(ctx)
	}

	// service opts
	srvOpts = append(srvOpts, micro.Name(Name))
	if i := time.Duration(ctx.Int("register_ttl")); i > 0 {
		srvOpts = append(srvOpts, micro.RegisterTTL(i*time.Second))
	}
	if i := time.Duration(ctx.Int("register_interval")); i > 0 {
		srvOpts = append(srvOpts, micro.RegisterInterval(i*time.Second))
	}

	// set the context
	var popts []proxy.Option

	// create new router
	var r router.Router

	routerName := ctx.String("router")
	routerAddr := ctx.String("router_address")

	ropts := []router.Option{
		router.Id(server.DefaultId),
		router.Client(client.DefaultClient),
		router.Address(routerAddr),
		router.Registry(registry.DefaultRegistry),
	}

	// check if we need to use the router service
	switch {
	case routerName == "go.micro.router":
		r = rs.NewRouter(ropts...)
	case routerName == "service":
		r = rs.NewRouter(ropts...)
	case len(routerAddr) > 0:
		r = rs.NewRouter(ropts...)
	default:
		r = router.NewRouter(ropts...)
	}

	// start the router
	if err := r.Start(); err != nil {
		log.Errorf("Proxy error starting router: %s", err)
		os.Exit(1)
	}

	popts = append(popts, proxy.WithRouter(r))

	// new proxy
	var p proxy.Proxy
	var srv server.Server

	// set endpoint
	if len(Endpoint) > 0 {
		switch {
		case strings.HasPrefix(Endpoint, "grpc://"):
			ep := strings.TrimPrefix(Endpoint, "grpc://")
			popts = append(popts, proxy.WithEndpoint(ep))
			p = grpc.NewProxy(popts...)
		case strings.HasPrefix(Endpoint, "http://"):
			// TODO: strip prefix?
			popts = append(popts, proxy.WithEndpoint(Endpoint))
			p = http.NewProxy(popts...)
		default:
			// TODO: strip prefix?
			popts = append(popts, proxy.WithEndpoint(Endpoint))
			p = mucp.NewProxy(popts...)
		}
	}

	serverOpts := []server.Option{
		server.Address(Address),
		server.Registry(rmem.NewRegistry()),
		server.Broker(bmem.NewBroker()),
	}

	// enable acme will create a net.Listener which
	if ctx.Bool("enable_acme") {
		var ap acme.Provider

		switch ACMEProvider {
		case "autocert":
			ap = autocert.NewProvider()
		case "certmagic":
			if ACMEChallengeProvider != "cloudflare" {
				log.Fatal("The only implemented DNS challenge provider is cloudflare")
			}
			apiToken, accountID := os.Getenv("CF_API_TOKEN"), os.Getenv("CF_ACCOUNT_ID")
			kvID := os.Getenv("KV_NAMESPACE_ID")
			if len(apiToken) == 0 || len(accountID) == 0 {
				log.Fatal("env variables CF_API_TOKEN and CF_ACCOUNT_ID must be set")
			}
			if len(kvID) == 0 {
				log.Fatal("env var KV_NAMESPACE_ID must be set to your cloudflare workers KV namespace ID")
			}

			cloudflareStore := cfstore.NewStore(
				cfstore.Token(apiToken),
				cfstore.Account(accountID),
				cfstore.Namespace(kvID),
			)
			storage := certmagic.NewStorage(
				memory.NewLock(),
				cloudflareStore,
			)
			config := cloudflare.NewDefaultConfig()
			config.AuthToken = apiToken
			config.ZoneToken = apiToken
			challengeProvider, err := cloudflare.NewDNSProviderConfig(config)
			if err != nil {
				log.Fatal(err.Error())
			}

			// define the provider
			ap = certmagic.NewProvider(
				acme.AcceptToS(true),
				acme.CA(ACMECA),
				acme.Cache(storage),
				acme.ChallengeProvider(challengeProvider),
				acme.OnDemand(false),
			)
		default:
			log.Fatalf("Unsupported acme provider: %s\n", ACMEProvider)
		}

		// generate the tls config
		config, err := ap.TLSConfig(helper.ACMEHosts(ctx)...)
		if err != nil {
			log.Fatalf("Failed to generate acme tls config: %v", err)
		}

		// set the tls config
		serverOpts = append(serverOpts, server.TLSConfig(config))
		// enable tls will leverage tls certs and generate a tls.Config
	} else if ctx.Bool("enable_tls") {
		// get certificates from the context
		config, err := helper.TLSConfig(ctx)
		if err != nil {
			log.Fatal(err)
			return
		}
		serverOpts = append(serverOpts, server.TLSConfig(config))
	}

	// set proxy
	if p == nil && len(Protocol) > 0 {
		switch Protocol {
		case "http":
			p = http.NewProxy(popts...)
			// TODO: http server
		case "mucp":
			popts = append(popts, proxy.WithClient(mucli.NewClient()))
			p = mucp.NewProxy(popts...)

			serverOpts = append(serverOpts, server.WithRouter(p))
			srv = server.NewServer(serverOpts...)
		default:
			p = mucp.NewProxy(popts...)

			serverOpts = append(serverOpts, server.WithRouter(p))
			srv = sgrpc.NewServer(serverOpts...)
		}
	}

	if len(Endpoint) > 0 {
		log.Infof("Proxy [%s] serving endpoint: %s", p.String(), Endpoint)
	} else {
		log.Infof("Proxy [%s] serving protocol: %s", p.String(), Protocol)
	}

	// new service
	service := micro.NewService(srvOpts...)

	// create a new proxy muxer which includes the debug handler
	muxer := mux.New(Name, p)

	// set the router
	service.Server().Init(
		server.WithRouter(muxer),
	)

	// Start the proxy server
	if err := srv.Start(); err != nil {
		log.Fatal(err)
	}

	// Run internal service
	if err := service.Run(); err != nil {
		log.Fatal(err)
	}

	// Stop the server
	if err := srv.Stop(); err != nil {
		log.Fatal(err)
	}
}

func Commands(options ...micro.Option) []*cli.Command {
	command := &cli.Command{
		Name:  "proxy",
		Usage: "Run the service proxy",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:    "router",
				Usage:   "Set the router to use e.g default, go.micro.router",
				EnvVars: []string{"MICRO_ROUTER"},
			},
			&cli.StringFlag{
				Name:    "router_address",
				Usage:   "Set the router address",
				EnvVars: []string{"MICRO_ROUTER_ADDRESS"},
			},
			&cli.StringFlag{
				Name:    "address",
				Usage:   "Set the proxy http address e.g 0.0.0.0:8081",
				EnvVars: []string{"MICRO_PROXY_ADDRESS"},
			},
			&cli.StringFlag{
				Name:    "protocol",
				Usage:   "Set the protocol used for proxying e.g mucp, grpc, http",
				EnvVars: []string{"MICRO_PROXY_PROTOCOL"},
			},
			&cli.StringFlag{
				Name:    "endpoint",
				Usage:   "Set the endpoint to route to e.g greeter or localhost:9090",
				EnvVars: []string{"MICRO_PROXY_ENDPOINT"},
			},
		},
		Action: func(ctx *cli.Context) error {
			run(ctx, options...)
			return nil
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []*cli.Command{command}
}
<commit_msg>Configure Proxy :443 Auth (#568)<commit_after>// Package proxy is a cli proxy
package proxy

import (
	"os"
	"strings"
	"time"

	"github.com/go-acme/lego/v3/providers/dns/cloudflare"
	"github.com/micro/cli/v2"
	"github.com/micro/go-micro/v2"
	"github.com/micro/go-micro/v2/api/server/acme"
	"github.com/micro/go-micro/v2/api/server/acme/autocert"
	"github.com/micro/go-micro/v2/api/server/acme/certmagic"
	"github.com/micro/go-micro/v2/auth"
	bmem "github.com/micro/go-micro/v2/broker/memory"
	"github.com/micro/go-micro/v2/client"
	mucli "github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/config/cmd"
	log "github.com/micro/go-micro/v2/logger"
	"github.com/micro/go-micro/v2/proxy"
	"github.com/micro/go-micro/v2/proxy/grpc"
	"github.com/micro/go-micro/v2/proxy/http"
	"github.com/micro/go-micro/v2/proxy/mucp"
	"github.com/micro/go-micro/v2/registry"
	rmem "github.com/micro/go-micro/v2/registry/memory"
	"github.com/micro/go-micro/v2/router"
	rs "github.com/micro/go-micro/v2/router/service"
	"github.com/micro/go-micro/v2/server"
	sgrpc "github.com/micro/go-micro/v2/server/grpc"
	cfstore "github.com/micro/go-micro/v2/store/cloudflare"
	"github.com/micro/go-micro/v2/sync/lock/memory"
	"github.com/micro/go-micro/v2/util/mux"
	"github.com/micro/go-micro/v2/util/wrapper"
	"github.com/micro/micro/v2/internal/helper"
)

var (
	// Name of the proxy
	Name = "go.micro.proxy"
	// The address of the proxy
	Address = ":8081"
	// the proxy protocol
	Protocol = "grpc"
	// The endpoint host to route to
	Endpoint string
	// ACME (Cert management)
	ACMEProvider          = "autocert"
	ACMEChallengeProvider = "cloudflare"
	ACMECA                = acme.LetsEncryptProductionCA
)

func run(ctx *cli.Context, srvOpts ...micro.Option) {
	log.Init(log.WithFields(map[string]interface{}{"service": "proxy"}))

	// because MICRO_PROXY_ADDRESS is used internally by the go-micro/client
	// we need to unset it so we don't end up calling ourselves infinitely
	os.Unsetenv("MICRO_PROXY_ADDRESS")

	if len(ctx.String("server_name")) > 0 {
		Name = ctx.String("server_name")
	}
	if len(ctx.String("address")) > 0 {
		Address = ctx.String("address")
	}
	if len(ctx.String("endpoint")) > 0 {
		Endpoint = ctx.String("endpoint")
	}
	if len(ctx.String("protocol")) > 0 {
		Protocol = ctx.String("protocol")
	}
	if len(ctx.String("acme_provider")) > 0 {
		ACMEProvider = ctx.String("acme_provider")
	}

	// Init plugins
	for _, p := range Plugins() {
		p.Init(ctx)
	}

	// service opts
	srvOpts = append(srvOpts, micro.Name(Name))
	if i := time.Duration(ctx.Int("register_ttl")); i > 0 {
		srvOpts = append(srvOpts, micro.RegisterTTL(i*time.Second))
	}
	if i := time.Duration(ctx.Int("register_interval")); i > 0 {
		srvOpts = append(srvOpts, micro.RegisterInterval(i*time.Second))
	}

	// set the context
	var popts []proxy.Option

	// create new router
	var r router.Router

	routerName := ctx.String("router")
	routerAddr := ctx.String("router_address")

	ropts := []router.Option{
		router.Id(server.DefaultId),
		router.Client(client.DefaultClient),
		router.Address(routerAddr),
		router.Registry(registry.DefaultRegistry),
	}

	// check if we need to use the router service
	switch {
	case routerName == "go.micro.router":
		r = rs.NewRouter(ropts...)
	case routerName == "service":
		r = rs.NewRouter(ropts...)
	case len(routerAddr) > 0:
		r = rs.NewRouter(ropts...)
	default:
		r = router.NewRouter(ropts...)
	}

	// start the router
	if err := r.Start(); err != nil {
		log.Errorf("Proxy error starting router: %s", err)
		os.Exit(1)
	}

	popts = append(popts, proxy.WithRouter(r))

	// new proxy
	var p proxy.Proxy
	var srv server.Server

	// set endpoint
	if len(Endpoint) > 0 {
		switch {
		case strings.HasPrefix(Endpoint, "grpc://"):
			ep := strings.TrimPrefix(Endpoint, "grpc://")
			popts = append(popts, proxy.WithEndpoint(ep))
			p = grpc.NewProxy(popts...)
		case strings.HasPrefix(Endpoint, "http://"):
			// TODO: strip prefix?
			popts = append(popts, proxy.WithEndpoint(Endpoint))
			p = http.NewProxy(popts...)
		default:
			// TODO: strip prefix?
			popts = append(popts, proxy.WithEndpoint(Endpoint))
			p = mucp.NewProxy(popts...)
		}
	}

	serverOpts := []server.Option{
		server.Address(Address),
		server.Registry(rmem.NewRegistry()),
		server.Broker(bmem.NewBroker()),
	}

	// enable acme will create a net.Listener which
	if ctx.Bool("enable_acme") {
		var ap acme.Provider

		switch ACMEProvider {
		case "autocert":
			ap = autocert.NewProvider()
		case "certmagic":
			if ACMEChallengeProvider != "cloudflare" {
				log.Fatal("The only implemented DNS challenge provider is cloudflare")
			}
			apiToken, accountID := os.Getenv("CF_API_TOKEN"), os.Getenv("CF_ACCOUNT_ID")
			kvID := os.Getenv("KV_NAMESPACE_ID")
			if len(apiToken) == 0 || len(accountID) == 0 {
				log.Fatal("env variables CF_API_TOKEN and CF_ACCOUNT_ID must be set")
			}
			if len(kvID) == 0 {
				log.Fatal("env var KV_NAMESPACE_ID must be set to your cloudflare workers KV namespace ID")
			}

			cloudflareStore := cfstore.NewStore(
				cfstore.Token(apiToken),
				cfstore.Account(accountID),
				cfstore.Namespace(kvID),
			)
			storage := certmagic.NewStorage(
				memory.NewLock(),
				cloudflareStore,
			)
			config := cloudflare.NewDefaultConfig()
			config.AuthToken = apiToken
			config.ZoneToken = apiToken
			challengeProvider, err := cloudflare.NewDNSProviderConfig(config)
			if err != nil {
				log.Fatal(err.Error())
			}

			// define the provider
			ap = certmagic.NewProvider(
				acme.AcceptToS(true),
				acme.CA(ACMECA),
				acme.Cache(storage),
				acme.ChallengeProvider(challengeProvider),
				acme.OnDemand(false),
			)
		default:
			log.Fatalf("Unsupported acme provider: %s\n", ACMEProvider)
		}

		// generate the tls config
		config, err := ap.TLSConfig(helper.ACMEHosts(ctx)...)
		if err != nil {
			log.Fatalf("Failed to generate acme tls config: %v", err)
		}

		// set the tls config
		serverOpts = append(serverOpts, server.TLSConfig(config))
		// enable tls will leverage tls certs and generate a tls.Config
	} else if ctx.Bool("enable_tls") {
		// get certificates from the context
		config, err := helper.TLSConfig(ctx)
		if err != nil {
			log.Fatal(err)
			return
		}
		serverOpts = append(serverOpts, server.TLSConfig(config))
	}

	// add auth wrapper to server
	if ctx.IsSet("auth") {
		a, ok := cmd.DefaultAuths[ctx.String("auth")]
		if !ok {
			log.Fatalf("%v is not a valid auth", ctx.String("auth"))
			return
		}

		var authOpts []auth.Option
		if ctx.IsSet("auth_exclude") {
			authOpts = append(authOpts, auth.Excludes(ctx.StringSlice("auth_exclude")...))
		}
		if ctx.IsSet("auth_public_key") {
			authOpts = append(authOpts, auth.PublicKey(ctx.String("auth_public_key")))
		}
		if ctx.IsSet("auth_private_key") {
			authOpts = append(authOpts, auth.PublicKey(ctx.String("auth_private_key")))
		}

		authFn := func() auth.Auth { return a(authOpts...) }
		authOpt := server.WrapHandler(wrapper.AuthHandler(authFn))
		serverOpts = append(serverOpts, authOpt)
	}

	// set proxy
	if p == nil && len(Protocol) > 0 {
		switch Protocol {
		case "http":
			p = http.NewProxy(popts...)
			// TODO: http server
		case "mucp":
			popts = append(popts, proxy.WithClient(mucli.NewClient()))
			p = mucp.NewProxy(popts...)

			serverOpts = append(serverOpts, server.WithRouter(p))
			srv = server.NewServer(serverOpts...)
		default:
			p = mucp.NewProxy(popts...)

			serverOpts = append(serverOpts, server.WithRouter(p))
			srv = sgrpc.NewServer(serverOpts...)
		}
	}

	if len(Endpoint) > 0 {
		log.Infof("Proxy [%s] serving endpoint: %s", p.String(), Endpoint)
	} else {
		log.Infof("Proxy [%s] serving protocol: %s", p.String(), Protocol)
	}

	// new service
	service := micro.NewService(srvOpts...)

	// create a new proxy muxer which includes the debug handler
	muxer := mux.New(Name, p)

	// set the router
	service.Server().Init(
		server.WithRouter(muxer),
	)

	// Start the proxy server
	if err := srv.Start(); err != nil {
		log.Fatal(err)
	}

	// Run internal service
	if err := service.Run(); err != nil {
		log.Fatal(err)
	}

	// Stop the server
	if err := srv.Stop(); err != nil {
		log.Fatal(err)
	}
}

func Commands(options ...micro.Option) []*cli.Command {
	command := &cli.Command{
		Name:  "proxy",
		Usage: "Run the service proxy",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:    "router",
				Usage:   "Set the router to use e.g default, go.micro.router",
				EnvVars: []string{"MICRO_ROUTER"},
			},
			&cli.StringFlag{
				Name:    "router_address",
				Usage:   "Set the router address",
				EnvVars: []string{"MICRO_ROUTER_ADDRESS"},
			},
			&cli.StringFlag{
				Name:    "address",
				Usage:   "Set the proxy http address e.g 0.0.0.0:8081",
				EnvVars: []string{"MICRO_PROXY_ADDRESS"},
			},
			&cli.StringFlag{
				Name:    "protocol",
				Usage:   "Set the protocol used for proxying e.g mucp, grpc, http",
				EnvVars: []string{"MICRO_PROXY_PROTOCOL"},
			},
			&cli.StringFlag{
				Name:    "endpoint",
				Usage:   "Set the endpoint to route to e.g greeter or localhost:9090",
				EnvVars: []string{"MICRO_PROXY_ENDPOINT"},
			},
			&cli.StringFlag{
				Name:    "auth",
				Usage:   "Set the proxy auth e.g jwt",
				EnvVars: []string{"MICRO_PROXY_AUTH"},
			},
		},
		Action: func(ctx *cli.Context) error {
			run(ctx, options...)
			return nil
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []*cli.Command{command}
}
<|endoftext|>


<commit_before>package delayd

import (
	"log"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/armon/consul-api"
	"github.com/hashicorp/raft"
	"github.com/streadway/amqp"
)

// a generous 60 seconds to apply raft commands
const raftMaxTime = time.Duration(60) * time.Second

// RoutingKey is used by AMQP consumer when binding a queue to an exchange.
const RoutingKey = "delayd"

// delaydService is used by consulapi to watch delayd service on Consul
const delaydService = "delayd"

// Server is the delayd server. It handles the server lifecycle (startup, clean shutdown)
type Server struct {
	sender       Sender
	receiver     Receiver
	raft         *Raft
	timer        *Timer
	shutdownCh   chan bool
	serviceCh    chan []*consulapi.ServiceEntry
	leader       bool
	mu           sync.Mutex
	consul       *consulapi.Client
	config       Config
	registration *consulapi.AgentServiceRegistration
	bootstrapped bool
}

// NewServer initialize Server instance.
func NewServer(c Config, sender Sender, receiver Receiver) (*Server, error) {
	if len(c.LogDir) != 0 {
		logFile := filepath.Join(c.LogDir, "delayd.log")
		logOutput, err := os.OpenFile(logFile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
		if err != nil {
			return nil, err
		}
		log.SetOutput(logOutput)
	}

	consul, err := NewConsul(c.Consul)
	if err != nil {
		return nil, err
	}

	raft, err := NewRaft(c.Raft, c.DataDir, c.LogDir)
	if err != nil {
		return nil, err
	}

	return &Server{
		sender:     sender,
		receiver:   receiver,
		raft:       raft,
		consul:     consul,
		config:     c,
		shutdownCh: make(chan bool),
		serviceCh:  make(chan []*consulapi.ServiceEntry),
	}, nil
}

// Run starts server and begins its main loop.
func (s *Server) Run() {
	Info("server: starting delayd")

	s.timer = NewTimer(s.timerSend)

	go s.observeLeaderChanges()
	go s.observeNextTime()

	if s.config.UseConsul {
		go s.startConsulBackend()
	}

	for {
		msg, ok := <-s.receiver.MessageCh()
		entry := msg.Entry
		// XXX cleanup needed here before exit
		if !ok {
			continue
		}

		Debug("server: got new request entry:", entry)
		b, err := entry.ToBytes()
		if err != nil {
			Error("server: error encoding entry: ", err)
			continue
		}

		if err := s.raft.Add(b, raftMaxTime); err != nil {
			Error("server: failed to add: ", err)
			msg.Nack()
			continue
		}

		msg.Ack()
	}
}

// Stop shuts down the Server cleanly. Order of the Close calls is important.
func (s *Server) Stop() {
	Info("server: shutting down gracefully.")

	// stop triggering new changes to the FSM
	s.receiver.Close()
	s.timer.Stop()

	if s.config.UseConsul {
		s.deregisterService()
	}

	close(s.shutdownCh)

	// with no FSM stimulus, we don't need to send.
	s.sender.Close()

	s.raft.Close()

	Info("server: terminated.")
}

func (s *Server) resetTimer() {
	ok, t, err := s.raft.fsm.store.NextTime()
	if err != nil {
		Fatal("server: could not read initial send time from storage:", err)
	}
	if ok {
		s.timer.Reset(t, true)
	}
}

func (s *Server) registerService() error {
	Info("server: registering delayd service on Consul")
	agent := s.consul.Agent()

	_, port, err := net.SplitHostPort(s.config.Raft.Listen)
	if err != nil {
		return err
	}
	p, err := strconv.ParseInt(port, 10, 32)
	if err != nil {
		return nil
	}

	s.registration = &consulapi.AgentServiceRegistration{
		Name: delaydService,
		ID:   "delayd-" + port,
		Port: int(p),
	}
	err = agent.ServiceRegister(s.registration)
	if err != nil {
		Warn("server: failed to register delayd service on Consul.", err)
	}
	return err
}

func (s *Server) deregisterService() error {
	Info("server: deregistering delayd service from Consul")
	agent := s.consul.Agent()
	err := agent.ServiceDeregister(s.registration.ID)
	if err != nil {
		Warn("server: failed to deregister delayd service from Consul")
	}
	return err
}

func (s *Server) startConsulBackend() {
	go s.observeService()
	go s.observeServiceChanges()

	Infof("server: Waiting for joining %d nodes...", s.config.BootstrapExpect)

	// Do our best for registering delayd service to Consul
	for {
		if err := s.registerService(); err != nil {
			// TODO: exponential backoff
			time.Sleep(500 * time.Millisecond)
			continue
		}
		break
	}
}

func (s *Server) observeService() {
	query := consulapi.QueryOptions{
		AllowStale:        false,
		RequireConsistent: true,
	}
	health := s.consul.Health()

	Debug("server: starting service monitoring")
	for {
		select {
		case <-s.shutdownCh:
			close(s.serviceCh)
			return
		default:
		}

		Info("server: waiting for node changes...")
		entries, meta, err := health.Service(delaydService, "", false, &query)
		if err != nil {
			Error(err)
			continue
		}
		s.serviceCh <- entries
		query.WaitIndex = meta.LastIndex
	}
}

func (s *Server) registeredSelf(services []*consulapi.ServiceEntry) bool {
	for _, service := range services {
		if s.registration.Name == service.Service.Service &&
			s.registration.ID == service.Service.ID {
			return true
		}
	}
	return false
}

func (s *Server) bootstrap(services []*consulapi.ServiceEntry) {
	peers := convertPeersFromService(services)
	s.raft.raft.SetPeers(peers)
	s.bootstrapped = true
	Infof("server: bootstrapped with %v", peers)
}

func (s *Server) observeServiceChanges() {
	for {
		select {
		case <-s.shutdownCh:
			return
		case services := <-s.serviceCh:
			if !s.bootstrapped {
				if len(services) >= s.config.BootstrapExpect && s.registeredSelf(services) {
					s.bootstrap(services)
				}
				continue
			}

			// leader need to maintain peers in response to service changes after bootstraped.
			if !s.leader {
				continue
			}

			newPeers := convertPeersFromService(services)
			curPeers, err := s.raft.peerStore.Peers()
			if err != nil {
				Error("failed to get peers:", err)
				continue
			}

			for _, p := range newPeers {
				if !raft.PeerContained(curPeers, p) {
					// new peer is found..
					Infof("server: new peer %v is found on consul. adding...", p)
					s.raft.raft.AddPeer(p)
				}
			}
		}
	}
}

// Listen for changes to raft for when we transition to and from the leader state
// and react accordingly.
func (s *Server) observeLeaderChanges() {
	for isLeader := range s.raft.LeaderCh() {
		s.mu.Lock()
		s.leader = isLeader
		s.mu.Unlock()

		if isLeader {
			Debug("server: became raft leader")
			s.resetTimer()
			if err := s.receiver.Start(); err != nil {
				Fatal("server: error while starting receiver:", err)
			}
		} else {
			Debug("server: lost raft leadership")
			s.timer.Pause()
			if err := s.receiver.Pause(); err != nil {
				Fatal("server: error while starting receiver:", err)
			}
		}
	}
}

// Listen to storage for next time changes on entry commits
func (s *Server) observeNextTime() {
	for {
		select {
		case <-s.shutdownCh:
			return
		case t := <-s.raft.fsm.store.C:
			s.mu.Lock()
			if s.leader {
				Debug("server: got time: ", t)
				s.timer.Reset(t, false)
			}
			s.mu.Unlock()
		}
	}
}

func (s *Server) timerSend(t time.Time) {
	// FIXME: In the case of error, we don't remove a log from raft log.
	uuids, entries, err := s.raft.fsm.store.Get(t)
	if err != nil {
		Fatal("server: could not read entries from db:", err)
	}

	if len(entries) > 1 {
		Infof("server: sending %d entries", len(entries))
	}

	if len(entries) == 0 {
		return
	}

	for i, e := range entries {
		if err := s.sender.Send(e); err != nil {
			if aerr, ok := err.(*amqp.Error); ok && aerr.Code == 504 {
				// error 504 code means that the exchange we were trying
				// to send on didnt exist.  In the case of delayd this usually
				// means that a consumer didn't set up the exchange they wish
				// to be notified on. We do not attempt to make this for them,
				// as we don't know what exchange options they would want, we
				// simply drop this message, other errors are fatal
				Warnf("server: channel/connection not set up for exchange `%s`, message will be deleted", e.Target, aerr)
			}

			// FIXME: I don't think Fatal here is right way.
			// If a reason that node is failed to send is node specific,
			// Fatal causes leader election then a problem may be resolved.
			// If the reason is not node specific, all instance may be down....
			Fatal("server: could not send entry:", err)
		}

		if err := s.raft.Remove(uuids[i], raftMaxTime); err != nil {
			// This node is no longer the leader. give up on other amqp sends,
			// and scheduling the next emission
			Warnf("server: lost raft leadership during remove. AMQP send will be a duplicate. uuid=%x", uuids[i])
			break
		}
	}

	// ensure everyone is up to date
	Debug("server: syncing raft after send.")
	err = s.raft.SyncAll()
	if err != nil {
		Warn("server: lost raft leadership during sync after send.")
		return
	}

	Debug("server: synced raft after send.")
}

func convertPeersFromService(services []*consulapi.ServiceEntry) []net.Addr {
	peers := []net.Addr{}
	for _, service := range services {
		port := strconv.FormatInt(int64(service.Service.Port), 10)
		// Currently, IP address in consul's service catalog can not be set in registration.
		// This limitation is so bad for testing multiple delayd instance with 1 agent.
		// See https://github.com/hashicorp/consul/issues/229 for discussion
		address := service.Node.Address
		if tweak := os.Getenv("DELAYD_CONSUL_AGENT_SERVICE_ADDRESS_TWEAK"); tweak != "" {
			address = tweak
		}
		peer, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(address, port))
		if err != nil {
			Fatal("server: bad peer:", err)
		}
		peers = append(peers, peer)
	}
	return peers
}
<commit_msg>server: Fix typo<commit_after>package delayd

import (
	"log"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/armon/consul-api"
	"github.com/hashicorp/raft"
	"github.com/streadway/amqp"
)

// a generous 60 seconds to apply raft commands
const raftMaxTime = time.Duration(60) * time.Second

// RoutingKey is used by AMQP consumer when binding a queue to an exchange.
const RoutingKey = "delayd"

// delaydService is used by consulapi to watch delayd service on Consul
const delaydService = "delayd"

// Server is the delayd server. It handles the server lifecycle (startup, clean shutdown)
type Server struct {
	sender       Sender
	receiver     Receiver
	raft         *Raft
	timer        *Timer
	shutdownCh   chan bool
	serviceCh    chan []*consulapi.ServiceEntry
	leader       bool
	mu           sync.Mutex
	consul       *consulapi.Client
	config       Config
	registration *consulapi.AgentServiceRegistration
	bootstrapped bool
}

// NewServer initialize Server instance.
func NewServer(c Config, sender Sender, receiver Receiver) (*Server, error) {
	if len(c.LogDir) != 0 {
		logFile := filepath.Join(c.LogDir, "delayd.log")
		logOutput, err := os.OpenFile(logFile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
		if err != nil {
			return nil, err
		}
		log.SetOutput(logOutput)
	}

	consul, err := NewConsul(c.Consul)
	if err != nil {
		return nil, err
	}

	raft, err := NewRaft(c.Raft, c.DataDir, c.LogDir)
	if err != nil {
		return nil, err
	}

	return &Server{
		sender:     sender,
		receiver:   receiver,
		raft:       raft,
		consul:     consul,
		config:     c,
		shutdownCh: make(chan bool),
		serviceCh:  make(chan []*consulapi.ServiceEntry),
	}, nil
}

// Run starts server and begins its main loop.
func (s *Server) Run() {
	Info("server: starting delayd")

	s.timer = NewTimer(s.timerSend)

	go s.observeLeaderChanges()
	go s.observeNextTime()

	if s.config.UseConsul {
		go s.startConsulBackend()
	}

	for {
		msg, ok := <-s.receiver.MessageCh()
		entry := msg.Entry
		// XXX cleanup needed here before exit
		if !ok {
			continue
		}

		Debug("server: got new request entry:", entry)
		b, err := entry.ToBytes()
		if err != nil {
			Error("server: error encoding entry: ", err)
			continue
		}

		if err := s.raft.Add(b, raftMaxTime); err != nil {
			Error("server: failed to add: ", err)
			msg.Nack()
			continue
		}

		msg.Ack()
	}
}

// Stop shuts down the Server cleanly. Order of the Close calls is important.
func (s *Server) Stop() {
	Info("server: shutting down gracefully.")

	// stop triggering new changes to the FSM
	s.receiver.Close()
	s.timer.Stop()

	if s.config.UseConsul {
		s.deregisterService()
	}

	close(s.shutdownCh)

	// with no FSM stimulus, we don't need to send.
	s.sender.Close()

	s.raft.Close()

	Info("server: terminated.")
}

func (s *Server) resetTimer() {
	ok, t, err := s.raft.fsm.store.NextTime()
	if err != nil {
		Fatal("server: could not read initial send time from storage:", err)
	}
	if ok {
		s.timer.Reset(t, true)
	}
}

func (s *Server) registerService() error {
	Info("server: registering delayd service on Consul")
	agent := s.consul.Agent()

	_, port, err := net.SplitHostPort(s.config.Raft.Listen)
	if err != nil {
		return err
	}
	p, err := strconv.ParseInt(port, 10, 32)
	if err != nil {
		return nil
	}

	s.registration = &consulapi.AgentServiceRegistration{
		Name: delaydService,
		ID:   "delayd-" + port,
		Port: int(p),
	}
	err = agent.ServiceRegister(s.registration)
	if err != nil {
		Warn("server: failed to register delayd service on Consul.", err)
	}
	return err
}

func (s *Server) deregisterService() error {
	Info("server: deregistering delayd service from Consul")
	agent := s.consul.Agent()
	err := agent.ServiceDeregister(s.registration.ID)
	if err != nil {
		Warn("server: failed to deregister delayd service from Consul")
	}
	return err
}

func (s *Server) startConsulBackend() {
	go s.observeService()
	go s.observeServiceChanges()

	Infof("server: waiting for joining %d nodes...", s.config.BootstrapExpect)

	// Do our best for registering delayd service to Consul
	for {
		if err := s.registerService(); err != nil {
			// TODO: exponential backoff
			time.Sleep(500 * time.Millisecond)
			continue
		}
		break
	}
}

func (s *Server) observeService() {
	query := consulapi.QueryOptions{
		AllowStale:        false,
		RequireConsistent: true,
	}
	health := s.consul.Health()

	Debug("server: starting service monitoring")
	for {
		select {
		case <-s.shutdownCh:
			close(s.serviceCh)
			return
		default:
		}

		Info("server: waiting for node changes...")
		entries, meta, err := health.Service(delaydService, "", false, &query)
		if err != nil {
			Error(err)
			continue
		}
		s.serviceCh <- entries
		query.WaitIndex = meta.LastIndex
	}
}

func (s *Server) registeredSelf(services []*consulapi.ServiceEntry) bool {
	for _, service := range services {
		if s.registration.Name == service.Service.Service &&
			s.registration.ID == service.Service.ID {
			return true
		}
	}
	return false
}

func (s *Server) bootstrap(services []*consulapi.ServiceEntry) {
	peers := convertPeersFromService(services)
	s.raft.raft.SetPeers(peers)
	s.bootstrapped = true
	Infof("server: bootstrapped with %v", peers)
}

func (s *Server) observeServiceChanges() {
	for {
		select {
		case <-s.shutdownCh:
			return
		case services := <-s.serviceCh:
			if !s.bootstrapped {
				if len(services) >= s.config.BootstrapExpect && s.registeredSelf(services) {
					s.bootstrap(services)
				}
				continue
			}

			// leader needs to maintain peers in response to service changes after bootstraped.
			if !s.leader {
				continue
			}

			newPeers := convertPeersFromService(services)
			curPeers, err := s.raft.peerStore.Peers()
			if err != nil {
				Error("failed to get peers:", err)
				continue
			}

			for _, p := range newPeers {
				if !raft.PeerContained(curPeers, p) {
					// new peer is found..
					Infof("server: new peer %v is found on consul. adding...", p)
					s.raft.raft.AddPeer(p)
				}
			}
		}
	}
}

// Listen for changes to raft for when we transition to and from the leader state
// and react accordingly.
func (s *Server) observeLeaderChanges() {
	for isLeader := range s.raft.LeaderCh() {
		s.mu.Lock()
		s.leader = isLeader
		s.mu.Unlock()

		if isLeader {
			Debug("server: became raft leader")
			s.resetTimer()
			if err := s.receiver.Start(); err != nil {
				Fatal("server: error while starting receiver:", err)
			}
		} else {
			Debug("server: lost raft leadership")
			s.timer.Pause()
			if err := s.receiver.Pause(); err != nil {
				Fatal("server: error while starting receiver:", err)
			}
		}
	}
}

// Listen to storage for next time changes on entry commits
func (s *Server) observeNextTime() {
	for {
		select {
		case <-s.shutdownCh:
			return
		case t := <-s.raft.fsm.store.C:
			s.mu.Lock()
			if s.leader {
				Debug("server: got time: ", t)
				s.timer.Reset(t, false)
			}
			s.mu.Unlock()
		}
	}
}

func (s *Server) timerSend(t time.Time) {
	// FIXME: In the case of error, we don't remove a log from raft log.
	uuids, entries, err := s.raft.fsm.store.Get(t)
	if err != nil {
		Fatal("server: could not read entries from db:", err)
	}

	if len(entries) > 1 {
		Infof("server: sending %d entries", len(entries))
	}

	if len(entries) == 0 {
		return
	}

	for i, e := range entries {
		if err := s.sender.Send(e); err != nil {
			if aerr, ok := err.(*amqp.Error); ok && aerr.Code == 504 {
				// error 504 code means that the exchange we were trying
				// to send on didnt exist.  In the case of delayd this usually
				// means that a consumer didn't set up the exchange they wish
				// to be notified on. We do not attempt to make this for them,
				// as we don't know what exchange options they would want, we
				// simply drop this message, other errors are fatal
				Warnf("server: channel/connection not set up for exchange `%s`, message will be deleted", e.Target, aerr)
			}

			// FIXME: I don't think Fatal here is right way.
			// If a reason that node is failed to send is node specific,
			// Fatal causes leader election then a problem may be resolved.
			// If the reason is not node specific, all instance may be down....
			Fatal("server: could not send entry:", err)
		}

		if err := s.raft.Remove(uuids[i], raftMaxTime); err != nil {
			// This node is no longer the leader. give up on other amqp sends,
			// and scheduling the next emission
			Warnf("server: lost raft leadership during remove. AMQP send will be a duplicate. uuid=%x", uuids[i])
			break
		}
	}

	// ensure everyone is up to date
	Debug("server: syncing raft after send.")
	err = s.raft.SyncAll()
	if err != nil {
		Warn("server: lost raft leadership during sync after send.")
		return
	}

	Debug("server: synced raft after send.")
}

func convertPeersFromService(services []*consulapi.ServiceEntry) []net.Addr {
	peers := []net.Addr{}
	for _, service := range services {
		port := strconv.FormatInt(int64(service.Service.Port), 10)
		// Currently, IP address in consul's service catalog can not be set in registration.
		// This limitation is so bad for testing multiple delayd instance with 1 agent.
		// See https://github.com/hashicorp/consul/issues/229 for discussion
		address := service.Node.Address
		if tweak := os.Getenv("DELAYD_CONSUL_AGENT_SERVICE_ADDRESS_TWEAK"); tweak != "" {
			address = tweak
		}
		peer, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(address, port))
		if err != nil {
			Fatal("server: bad peer:", err)
		}
		peers = append(peers, peer)
	}
	return peers
}
<|endoftext|>


<commit_before>package proxy

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"reflect"
	"regexp"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/andystanton/proxybastard/util"
	"github.com/deckarep/golang-set"
)

func addToMap(frequencyMap map[string]int, value string) map[string]int {
	if len(strings.TrimSpace(value)) > 0 {
		if _, ok := frequencyMap[value]; !ok {
			frequencyMap[value] = 0
		}
		frequencyMap[value] = frequencyMap[value] + 1
	}
	return frequencyMap
}

func awaitInput(prompt string, pattern string, prefix string) string {
	var matched string
	var found bool
	fmt.Printf("%s%s\n", prefix, prompt)
	fmt.Printf("%s> ", prefix)
	for i := 0; i < 3; i++ {
		reader := bufio.NewReader(os.Stdin)
		text, _ := reader.ReadString('\n')
		rPattern := regexp.MustCompile(pattern)

		if result := strings.TrimSpace(text); rPattern.MatchString(result) {
			matched = result
			found = true
			break
		} else if i < 2 {
			log.Printf("%sThat doesn't look right - try again...\n", prefix)
			log.Println()
			fmt.Printf("%s%s\n", prefix, prompt)
			fmt.Printf("%s> ", prefix)
		} else {
			log.Println()
			log.Print("Three failed attempts - aborting!")
		}
	}
	if !found {
		os.Exit(1)
	}
	return matched
}

func awaitFileInput(prompt string, prefix string) string {
	var matched string
	var found bool
	fmt.Printf("%s%s\n", prefix, prompt)
	fmt.Printf("%s> ", prefix)
	for i := 0; i < 3; i++ {
		reader := bufio.NewReader(os.Stdin)
		text, _ := reader.ReadString('\n')
		if file := util.SanitisePath(strings.TrimSpace(text)); util.FileExists(file) {
			matched = file
			found = true
			break
		} else if i < 2 {
			log.Printf("%sThat doesn't look right - try again...\n", prefix)
			log.Println()
			fmt.Printf("%s%s\n", prefix, prompt)
			fmt.Printf("%s> ", prefix)
		} else {
			log.Println()
			log.Print("Three failed attempts - aborting!")
		}
	}
	if !found {
		os.Exit(1)
	}
	return matched
}

func awaitYesInput(message string, prefix string) bool {
	input := awaitInput(message, "(y|n|^$)", prefix)
	return strings.EqualFold(input, "y") || strings.EqualFold(input, "")
}

func awaitNoInput(message string, prefix string) bool {
	input := awaitInput(message, "(y|n|^$)", prefix)
	return strings.EqualFold(input, "n") || strings.EqualFold(input, "")
}

// Setup presents the user with setup options.
func Setup(version string, acceptDefaults bool) {
	suggestedConfiguration := suggestConfiguration()
	actualConfiguration := Configuration{}
	actualConfiguration.Version = version
	readyToWrite := false

	if acceptDefaults {
		actualConfiguration = suggestedConfiguration
		readyToWrite = true
	} else {
		httpProxySet := false
		if len(suggestedConfiguration.ProxyHost) > 0 {
			message := fmt.Sprintf("Use suggested http proxy %s:%s? [Yn]", suggestedConfiguration.ProxyHost, suggestedConfiguration.ProxyPort)

			if httpProxySet = awaitYesInput(message, ""); httpProxySet {
				actualConfiguration.ProxyHost = suggestedConfiguration.ProxyHost
				actualConfiguration.ProxyPort = suggestedConfiguration.ProxyPort
			}
			fmt.Println()
		}

		if !httpProxySet {
			proxyHostPattern := "(?:https?://)?(.+):(\\d+)"
			proxyHostRegexp := regexp.MustCompile(proxyHostPattern)

			message := "Please enter an http proxy e.g. http://proxybastard:1234"
			matches := proxyHostRegexp.FindStringSubmatch(awaitInput(message, proxyHostPattern, ""))

			actualConfiguration.ProxyHost = fmt.Sprintf("http://%s", matches[1])
			actualConfiguration.ProxyPort = matches[2]
			httpProxySet = true
			fmt.Println()
		}

		socksProxySet := false
		if len(suggestedConfiguration.SOCKSProxyHost) > 0 {
			message := fmt.Sprintf("Use suggested SOCKS proxy %s:%s? [Yn]", suggestedConfiguration.SOCKSProxyHost, suggestedConfiguration.SOCKSProxyPort)

			if socksProxySet = awaitYesInput(message, ""); socksProxySet {
				actualConfiguration.SOCKSProxyHost = suggestedConfiguration.SOCKSProxyHost
				actualConfiguration.SOCKSProxyPort = suggestedConfiguration.SOCKSProxyPort
			}
			fmt.Println()
		}

		if !socksProxySet {
			socksHostPattern := "(?:(.+):(\\d+)|^$)"
			sockHostRegexp := regexp.MustCompile(socksHostPattern)
			message := "Please enter a SOCKS proxy or press return for none e.g. socks.proxybastard:1234"
			matches := sockHostRegexp.FindStringSubmatch(awaitInput(message, socksHostPattern, ""))

			if len(matches) > 0 {
				socksProxySet = true
				actualConfiguration.SOCKSProxyHost = matches[1]
				actualConfiguration.SOCKSProxyPort = matches[2]
			}
			fmt.Println()
		}

		if suggestedConfiguration.Targets != nil {
			targetsField := reflect.Indirect(reflect.ValueOf(suggestedConfiguration.Targets))

			for i := 0; i < targetsField.NumField(); i++ {
				fieldName := targetsField.Type().Field(i).Name

				if !util.InterfaceIsZero(targetsField.Field(i).Interface()) {
					targetField := reflect.Indirect(reflect.ValueOf(targetsField.Field(i).Interface()))

					configurationSet := awaitYesInput(fmt.Sprintf("Found %s! Use suggested configuration? [Yn]", fieldName), "")

					if configurationSet {
						if actualConfiguration.Targets == nil {
							actualConfiguration.Targets = &TargetsConfiguration{}
						}
						actualField := reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets)).FieldByName(fieldName)
						actualField.Set(reflect.ValueOf(targetsField.Field(i).Interface()))
					} else {
						fmt.Println()

						if awaitYesInput(fmt.Sprintf("Enable %s? [Yn]", fieldName), "  ") {
							fmt.Println()

							if actualConfiguration.Targets == nil {
								actualConfiguration.Targets = &TargetsConfiguration{}
							}
							actualField := reflect.New(reflect.TypeOf(targetField.Interface())).Interface()
							actualFieldValue := reflect.ValueOf(actualField)
							actualFieldValuePtr := reflect.Indirect(actualFieldValue)

							actualFieldValuePtr.FieldByName("Enabled").Set(reflect.ValueOf(true))

							if util.ValueHasField(targetField, "Files") {
								fieldFiles := targetField.FieldByName("Files").Interface().([]string)

								if awaitYesInput(fmt.Sprintf("Use suggested file %s? [Yn]", strings.Join(fieldFiles, ", ")), "  ") {
									actualFieldValuePtr.FieldByName("Files").Set(reflect.ValueOf(fieldFiles))
								} else {
									fmt.Println()
									message := fmt.Sprintf("Enter the path to a valid %s configuration file", fieldName)
									input := awaitFileInput(message, "  ")
									actualFieldValuePtr.FieldByName("Files").Set(reflect.ValueOf([]string{input}))
								}
							}

							if util.ValueHasMethod(targetField, "CustomPrompt") {
								customMethod := targetField.MethodByName("CustomPrompt")
								actualField = customMethod.Call([]reflect.Value{})[0].Interface()
							}

							reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets)).FieldByName(fieldName).Set(actualFieldValue)
						}
					}
					fmt.Println()
				}
			}
		}

		w := tabwriter.NewWriter(os.Stdout, 20, 0, 1, ' ', 0)

		fmt.Fprintln(w, "Settings")
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w, "HTTP Proxy")
		fmt.Fprintf(w, " - Host\t : %s\n", actualConfiguration.ProxyHost)
		fmt.Fprintf(w, " - Port\t : %s\n", actualConfiguration.ProxyPort)

		if len(actualConfiguration.SOCKSProxyHost) > 0 {
			fmt.Fprintln(w)
			fmt.Fprintln(w, "SOCKS Proxy")
			fmt.Fprintf(w, " - Host\t : %s\n", actualConfiguration.SOCKSProxyHost)
			fmt.Fprintf(w, " - Port\t : %s\n", actualConfiguration.SOCKSProxyPort)
		}
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w)

		if actualConfiguration.Targets != nil {
			targetsField := reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets))

			for i := 0; i < targetsField.NumField(); i++ {
				fieldName := targetsField.Type().Field(i).Name

				if !util.InterfaceIsZero(targetsField.Field(i).Interface()) {
					targetField := reflect.Indirect(reflect.ValueOf(targetsField.Field(i).Interface()))

					fmt.Fprintln(w, fieldName)
					for j := 0; j < targetField.NumField(); j++ {
						fmt.Fprintf(w, " - %s\t : %v\n", targetField.Type().Field(j).Name, targetField.Field(j).Interface())
					}
					fmt.Fprintln(w)
				}
			}
		}
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w)
		w.Flush()

		input := awaitInput("Write these settings to ~/.proxybastard/config.json? [Yn]", "(y|n|^$)", "")
		fmt.Println()
		readyToWrite = strings.EqualFold(input, "y") || strings.EqualFold(input, "")
	}

	marshalled, err := json.MarshalIndent(actualConfiguration, "", "    ")
	if err != nil {
		log.Fatal(err)
	}
	// fmt.Printf("%s\n\n", string(marshalled))
	if readyToWrite {
		util.WriteSliceToFile(util.SanitisePath("~/.proxybastard/config.json"), strings.Split(string(marshalled), "\n"))
		fmt.Println("Done")
	} else {
		fmt.Println("kthx")
	}
}

func getHighestFrequency(frequencyMap map[string]int) string {
	var highestValue int
	var mostFrequentKey string

	keys := []string{}
	for key := range frequencyMap {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, key := range keys {
		if value := frequencyMap[key]; value > highestValue {
			highestValue = value
			mostFrequentKey = key
		}
	}
	return mostFrequentKey
}

func suggestConfiguration() Configuration {
	var suggestedConfiguration Configuration

	suggestedProxyHosts := make(map[string]int)
	suggestedProxyPorts := make(map[string]int)
	suggestedSOCKSProxyHosts := make(map[string]int)
	suggestedSOCKSProxyPorts := make(map[string]int)
	suggestedNonProxyHosts := mapset.NewSet()

	targetsField := reflect.Indirect(reflect.ValueOf(&TargetsConfiguration{}))
	for i := 0; i < targetsField.NumField(); i++ {

		configurationField := reflect.New(reflect.TypeOf(targetsField.Field(i).Interface()).Elem()).Interface()
		withConfig, hasConfig := configurationField.(WithConfig)
		if hasConfig {
			fieldName := targetsField.Type().Field(i).Name

			if suggestedItemConfiguration := withConfig.suggestConfiguration(); suggestedItemConfiguration != nil {
				if suggestedConfiguration.Targets == nil {
					suggestedConfiguration.Targets = &TargetsConfiguration{}
				}

				addToMap(suggestedProxyHosts, util.SanitiseHTTPProxyURL(suggestedItemConfiguration.ProxyHost))
				addToMap(suggestedProxyPorts, suggestedItemConfiguration.ProxyPort)
				addToMap(suggestedSOCKSProxyHosts, suggestedItemConfiguration.SOCKSProxyHost)
				addToMap(suggestedSOCKSProxyPorts, suggestedItemConfiguration.SOCKSProxyPort)

				for _, nonProxyHost := range suggestedItemConfiguration.NonProxyHosts {
					suggestedNonProxyHosts.Add(nonProxyHost)
				}

				targetsField := reflect.Indirect(reflect.ValueOf(suggestedConfiguration.Targets))
				targetsFieldSuggested := reflect.Indirect(reflect.ValueOf(suggestedItemConfiguration.Targets))

				targetsField.FieldByName(fieldName).Set(targetsFieldSuggested.FieldByName(fieldName))
			}
		}
	}

	suggestedConfiguration.ProxyHost = getHighestFrequency(suggestedProxyHosts)
	suggestedConfiguration.ProxyPort = getHighestFrequency(suggestedProxyPorts)
	suggestedConfiguration.SOCKSProxyHost = getHighestFrequency(suggestedSOCKSProxyHosts)
	suggestedConfiguration.SOCKSProxyPort = getHighestFrequency(suggestedSOCKSProxyPorts)
	for suggestedNonProxyHost := range suggestedNonProxyHosts.Iter() {
		suggestedConfiguration.NonProxyHosts = append(suggestedConfiguration.NonProxyHosts, suggestedNonProxyHost.(string))
	}

	return suggestedConfiguration
}
<commit_msg>#46 make ~/.proxybastard if it doesn't exist when writing config<commit_after>package proxy

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"reflect"
	"regexp"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/andystanton/proxybastard/util"
	"github.com/deckarep/golang-set"
)

func addToMap(frequencyMap map[string]int, value string) map[string]int {
	if len(strings.TrimSpace(value)) > 0 {
		if _, ok := frequencyMap[value]; !ok {
			frequencyMap[value] = 0
		}
		frequencyMap[value] = frequencyMap[value] + 1
	}
	return frequencyMap
}

func awaitInput(prompt string, pattern string, prefix string) string {
	var matched string
	var found bool
	fmt.Printf("%s%s\n", prefix, prompt)
	fmt.Printf("%s> ", prefix)
	for i := 0; i < 3; i++ {
		reader := bufio.NewReader(os.Stdin)
		text, _ := reader.ReadString('\n')
		rPattern := regexp.MustCompile(pattern)

		if result := strings.TrimSpace(text); rPattern.MatchString(result) {
			matched = result
			found = true
			break
		} else if i < 2 {
			log.Printf("%sThat doesn't look right - try again...\n", prefix)
			log.Println()
			fmt.Printf("%s%s\n", prefix, prompt)
			fmt.Printf("%s> ", prefix)
		} else {
			log.Println()
			log.Print("Three failed attempts - aborting!")
		}
	}
	if !found {
		os.Exit(1)
	}
	return matched
}

func awaitFileInput(prompt string, prefix string) string {
	var matched string
	var found bool
	fmt.Printf("%s%s\n", prefix, prompt)
	fmt.Printf("%s> ", prefix)
	for i := 0; i < 3; i++ {
		reader := bufio.NewReader(os.Stdin)
		text, _ := reader.ReadString('\n')
		if file := util.SanitisePath(strings.TrimSpace(text)); util.FileExists(file) {
			matched = file
			found = true
			break
		} else if i < 2 {
			log.Printf("%sThat doesn't look right - try again...\n", prefix)
			log.Println()
			fmt.Printf("%s%s\n", prefix, prompt)
			fmt.Printf("%s> ", prefix)
		} else {
			log.Println()
			log.Print("Three failed attempts - aborting!")
		}
	}
	if !found {
		os.Exit(1)
	}
	return matched
}

func awaitYesInput(message string, prefix string) bool {
	input := awaitInput(message, "(y|n|^$)", prefix)
	return strings.EqualFold(input, "y") || strings.EqualFold(input, "")
}

func awaitNoInput(message string, prefix string) bool {
	input := awaitInput(message, "(y|n|^$)", prefix)
	return strings.EqualFold(input, "n") || strings.EqualFold(input, "")
}

// Setup presents the user with setup options.
func Setup(version string, acceptDefaults bool) {
	suggestedConfiguration := suggestConfiguration()
	actualConfiguration := Configuration{}
	actualConfiguration.Version = version
	readyToWrite := false

	if acceptDefaults {
		actualConfiguration = suggestedConfiguration
		readyToWrite = true
	} else {
		httpProxySet := false
		if len(suggestedConfiguration.ProxyHost) > 0 {
			message := fmt.Sprintf("Use suggested http proxy %s:%s? [Yn]", suggestedConfiguration.ProxyHost, suggestedConfiguration.ProxyPort)

			if httpProxySet = awaitYesInput(message, ""); httpProxySet {
				actualConfiguration.ProxyHost = suggestedConfiguration.ProxyHost
				actualConfiguration.ProxyPort = suggestedConfiguration.ProxyPort
			}
			fmt.Println()
		}

		if !httpProxySet {
			proxyHostPattern := "(?:https?://)?(.+):(\\d+)"
			proxyHostRegexp := regexp.MustCompile(proxyHostPattern)

			message := "Please enter an http proxy e.g. http://proxybastard:1234"
			matches := proxyHostRegexp.FindStringSubmatch(awaitInput(message, proxyHostPattern, ""))

			actualConfiguration.ProxyHost = fmt.Sprintf("http://%s", matches[1])
			actualConfiguration.ProxyPort = matches[2]
			httpProxySet = true
			fmt.Println()
		}

		socksProxySet := false
		if len(suggestedConfiguration.SOCKSProxyHost) > 0 {
			message := fmt.Sprintf("Use suggested SOCKS proxy %s:%s? [Yn]", suggestedConfiguration.SOCKSProxyHost, suggestedConfiguration.SOCKSProxyPort)

			if socksProxySet = awaitYesInput(message, ""); socksProxySet {
				actualConfiguration.SOCKSProxyHost = suggestedConfiguration.SOCKSProxyHost
				actualConfiguration.SOCKSProxyPort = suggestedConfiguration.SOCKSProxyPort
			}
			fmt.Println()
		}

		if !socksProxySet {
			socksHostPattern := "(?:(.+):(\\d+)|^$)"
			sockHostRegexp := regexp.MustCompile(socksHostPattern)
			message := "Please enter a SOCKS proxy or press return for none e.g. socks.proxybastard:1234"
			matches := sockHostRegexp.FindStringSubmatch(awaitInput(message, socksHostPattern, ""))

			if len(matches) > 0 {
				socksProxySet = true
				actualConfiguration.SOCKSProxyHost = matches[1]
				actualConfiguration.SOCKSProxyPort = matches[2]
			}
			fmt.Println()
		}

		if suggestedConfiguration.Targets != nil {
			targetsField := reflect.Indirect(reflect.ValueOf(suggestedConfiguration.Targets))

			for i := 0; i < targetsField.NumField(); i++ {
				fieldName := targetsField.Type().Field(i).Name

				if !util.InterfaceIsZero(targetsField.Field(i).Interface()) {
					targetField := reflect.Indirect(reflect.ValueOf(targetsField.Field(i).Interface()))

					configurationSet := awaitYesInput(fmt.Sprintf("Found %s! Use suggested configuration? [Yn]", fieldName), "")

					if configurationSet {
						if actualConfiguration.Targets == nil {
							actualConfiguration.Targets = &TargetsConfiguration{}
						}
						actualField := reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets)).FieldByName(fieldName)
						actualField.Set(reflect.ValueOf(targetsField.Field(i).Interface()))
					} else {
						fmt.Println()

						if awaitYesInput(fmt.Sprintf("Enable %s? [Yn]", fieldName), "  ") {
							fmt.Println()

							if actualConfiguration.Targets == nil {
								actualConfiguration.Targets = &TargetsConfiguration{}
							}
							actualField := reflect.New(reflect.TypeOf(targetField.Interface())).Interface()
							actualFieldValue := reflect.ValueOf(actualField)
							actualFieldValuePtr := reflect.Indirect(actualFieldValue)

							actualFieldValuePtr.FieldByName("Enabled").Set(reflect.ValueOf(true))

							if util.ValueHasField(targetField, "Files") {
								fieldFiles := targetField.FieldByName("Files").Interface().([]string)

								if awaitYesInput(fmt.Sprintf("Use suggested file %s? [Yn]", strings.Join(fieldFiles, ", ")), "  ") {
									actualFieldValuePtr.FieldByName("Files").Set(reflect.ValueOf(fieldFiles))
								} else {
									fmt.Println()
									message := fmt.Sprintf("Enter the path to a valid %s configuration file", fieldName)
									input := awaitFileInput(message, "  ")
									actualFieldValuePtr.FieldByName("Files").Set(reflect.ValueOf([]string{input}))
								}
							}

							if util.ValueHasMethod(targetField, "CustomPrompt") {
								customMethod := targetField.MethodByName("CustomPrompt")
								actualField = customMethod.Call([]reflect.Value{})[0].Interface()
							}

							reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets)).FieldByName(fieldName).Set(actualFieldValue)
						}
					}
					fmt.Println()
				}
			}
		}

		w := tabwriter.NewWriter(os.Stdout, 20, 0, 1, ' ', 0)

		fmt.Fprintln(w, "Settings")
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w, "HTTP Proxy")
		fmt.Fprintf(w, " - Host\t : %s\n", actualConfiguration.ProxyHost)
		fmt.Fprintf(w, " - Port\t : %s\n", actualConfiguration.ProxyPort)

		if len(actualConfiguration.SOCKSProxyHost) > 0 {
			fmt.Fprintln(w)
			fmt.Fprintln(w, "SOCKS Proxy")
			fmt.Fprintf(w, " - Host\t : %s\n", actualConfiguration.SOCKSProxyHost)
			fmt.Fprintf(w, " - Port\t : %s\n", actualConfiguration.SOCKSProxyPort)
		}
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w)

		if actualConfiguration.Targets != nil {
			targetsField := reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets))

			for i := 0; i < targetsField.NumField(); i++ {
				fieldName := targetsField.Type().Field(i).Name

				if !util.InterfaceIsZero(targetsField.Field(i).Interface()) {
					targetField := reflect.Indirect(reflect.ValueOf(targetsField.Field(i).Interface()))

					fmt.Fprintln(w, fieldName)
					for j := 0; j < targetField.NumField(); j++ {
						fmt.Fprintf(w, " - %s\t : %v\n", targetField.Type().Field(j).Name, targetField.Field(j).Interface())
					}
					fmt.Fprintln(w)
				}
			}
		}
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w)
		w.Flush()

		input := awaitInput("Write these settings to ~/.proxybastard/config.json? [Yn]", "(y|n|^$)", "")
		fmt.Println()
		readyToWrite = strings.EqualFold(input, "y") || strings.EqualFold(input, "")
	}

	marshalled, err := json.MarshalIndent(actualConfiguration, "", "    ")
	if err != nil {
		log.Fatal(err)
	}

	if readyToWrite {
		s, err := os.Stat(util.SanitisePath("~/.proxybastard"))
		if err != nil && os.IsNotExist(err) {
			os.Mkdir(util.SanitisePath("~/.proxybastard"), 0755)
		} else if !s.IsDir() {
			panic("Unable to write to ~/.proxybastard")
		}
		// fmt.Printf("%s\n\n", string(marshalled))
		util.WriteSliceToFile(util.SanitisePath("~/.proxybastard/config.json"), strings.Split(string(marshalled), "\n"))
		fmt.Println("Wrote ~/.proxybastard/config.json")
	} else {
		fmt.Println("Aborting setup")
	}
}

func getHighestFrequency(frequencyMap map[string]int) string {
	var highestValue int
	var mostFrequentKey string

	keys := []string{}
	for key := range frequencyMap {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, key := range keys {
		if value := frequencyMap[key]; value > highestValue {
			highestValue = value
			mostFrequentKey = key
		}
	}
	return mostFrequentKey
}

func suggestConfiguration() Configuration {
	var suggestedConfiguration Configuration

	suggestedProxyHosts := make(map[string]int)
	suggestedProxyPorts := make(map[string]int)
	suggestedSOCKSProxyHosts := make(map[string]int)
	suggestedSOCKSProxyPorts := make(map[string]int)
	suggestedNonProxyHosts := mapset.NewSet()

	targetsField := reflect.Indirect(reflect.ValueOf(&TargetsConfiguration{}))
	for i := 0; i < targetsField.NumField(); i++ {

		configurationField := reflect.New(reflect.TypeOf(targetsField.Field(i).Interface()).Elem()).Interface()
		withConfig, hasConfig := configurationField.(WithConfig)
		if hasConfig {
			fieldName := targetsField.Type().Field(i).Name

			if suggestedItemConfiguration := withConfig.suggestConfiguration(); suggestedItemConfiguration != nil {
				if suggestedConfiguration.Targets == nil {
					suggestedConfiguration.Targets = &TargetsConfiguration{}
				}

				addToMap(suggestedProxyHosts, util.SanitiseHTTPProxyURL(suggestedItemConfiguration.ProxyHost))
				addToMap(suggestedProxyPorts, suggestedItemConfiguration.ProxyPort)
				addToMap(suggestedSOCKSProxyHosts, suggestedItemConfiguration.SOCKSProxyHost)
				addToMap(suggestedSOCKSProxyPorts, suggestedItemConfiguration.SOCKSProxyPort)

				for _, nonProxyHost := range suggestedItemConfiguration.NonProxyHosts {
					suggestedNonProxyHosts.Add(nonProxyHost)
				}

				targetsField := reflect.Indirect(reflect.ValueOf(suggestedConfiguration.Targets))
				targetsFieldSuggested := reflect.Indirect(reflect.ValueOf(suggestedItemConfiguration.Targets))

				targetsField.FieldByName(fieldName).Set(targetsFieldSuggested.FieldByName(fieldName))
			}
		}
	}

	suggestedConfiguration.ProxyHost = getHighestFrequency(suggestedProxyHosts)
	suggestedConfiguration.ProxyPort = getHighestFrequency(suggestedProxyPorts)
	suggestedConfiguration.SOCKSProxyHost = getHighestFrequency(suggestedSOCKSProxyHosts)
	suggestedConfiguration.SOCKSProxyPort = getHighestFrequency(suggestedSOCKSProxyPorts)
	for suggestedNonProxyHost := range suggestedNonProxyHosts.Iter() {
		suggestedConfiguration.NonProxyHosts = append(suggestedConfiguration.NonProxyHosts, suggestedNonProxyHost.(string))
	}

	return suggestedConfiguration
}
<|endoftext|>