<commit_before>package main

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"errors"
	"github.com/gorilla/mux"
	"log"
	"net/http"
	"time"
)

const (
	pubLen int = 12
)

func checkOTP(w http.ResponseWriter, r *http.Request, dal *Dal) {
	if r.URL.Query()["otp"] == nil || r.URL.Query()["nonce"] == nil || r.URL.Query()["id"] == nil {
		reply(w, "", "", "", MISSING_PARAMETER, "", dal)
		return
	}
	otp := r.URL.Query()["otp"][0]
	nonce := r.URL.Query()["nonce"][0]
	id := r.URL.Query()["id"][0]
	name := ""

	if len(otp) < pubLen {
		reply(w, otp, name, nonce, BAD_OTP, id, dal)
		return
	}
	pub := otp[:pubLen]

	k, err := dal.GetKey(pub)
	if err != nil {
		reply(w, otp, name, nonce, BAD_OTP, id, dal)
		return
	} else {
		k, err = Gate(k, otp)
		if err != nil {
			reply(w, otp, name, nonce, err.Error(), id, dal)
			return
		} else {
			err = dal.UpdateKey(k)
			if err != nil {
				log.Println("fail to update key counter/session")
				return
			}

			name := k.Name
			reply(w, otp, name, nonce, OK, id, dal)
			return
		}
	}
}

func Sign(values []string, key []byte) []byte {
	payload := ""
	for _, v := range values {
		payload += v + "&"
	}
	payload = payload[:len(payload)-1]

	mac := hmac.New(sha1.New, key)
	mac.Write([]byte(payload))
	return mac.Sum(nil)
}

func loadKey(id string, dal *Dal) ([]byte, error) {
	i, err := dal.GetApp(id)
	if err != nil {
		return []byte{}, errors.New(NO_SUCH_CLIENT)
	}

	return i, nil
}

func reply(w http.ResponseWriter, otp, name, nonce, status, id string, dal *Dal) {
	values := []string{}
	key := []byte{}
	err := errors.New("")

	values = append(values, "nonce="+nonce)
	values = append(values, "otp="+otp)
	if status != MISSING_PARAMETER {
		key, err = loadKey(id, dal)
		if err == nil {
			values = append(values, "name="+name)
			values = append(values, "status="+status)
		} else {
			values = append(values, "status="+err.Error())
		}
	} else {
		values = append(values, "status="+status)
	}
	values = append(values, "t="+time.Now().Format(time.RFC3339))
	if status != MISSING_PARAMETER {
		values = append(values, "h="+base64.StdEncoding.EncodeToString(Sign(values, key)))
	}

	ret := ""
	for _, v := range values {
		ret += v + "\n"
	}

	w.Write([]byte(ret))
}

func runAPI(dal *Dal, host, port string) {
	r := mux.NewRouter()

	r.HandleFunc("/wsapi/2.0/verify", func(w http.ResponseWriter, r *http.Request) {
		checkOTP(w, r, dal)
	}).Methods("GET")

	http.Handle("/", r)
	log.Printf("Listening on: %s:%s...", host, port)
	http.ListenAndServe(host+":"+port, nil)
}
<commit_msg>HMAC signature should be generated from sorted keys<commit_after>package main

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"errors"
	"github.com/gorilla/mux"
	"log"
	"net/http"
	"time"
	"sort"
)

const (
	pubLen int = 12
)

func checkOTP(w http.ResponseWriter, r *http.Request, dal *Dal) {
	if r.URL.Query()["otp"] == nil || r.URL.Query()["nonce"] == nil || r.URL.Query()["id"] == nil {
		reply(w, "", "", "", MISSING_PARAMETER, "", dal)
		return
	}
	otp := r.URL.Query()["otp"][0]
	nonce := r.URL.Query()["nonce"][0]
	id := r.URL.Query()["id"][0]
	name := ""

	if len(otp) < pubLen {
		reply(w, otp, name, nonce, BAD_OTP, id, dal)
		return
	}
	pub := otp[:pubLen]

	k, err := dal.GetKey(pub)
	if err != nil {
		reply(w, otp, name, nonce, BAD_OTP, id, dal)
		return
	} else {
		k, err = Gate(k, otp)
		if err != nil {
			reply(w, otp, name, nonce, err.Error(), id, dal)
			return
		} else {
			err = dal.UpdateKey(k)
			if err != nil {
				log.Println("fail to update key counter/session")
				return
			}

			name := k.Name
			reply(w, otp, name, nonce, OK, id, dal)
			return
		}
	}
}

func Sign(values []string, key []byte) []byte {
	payload := ""
        // Alphabetically sort the set of key/value pairs by key order.
        // ref: https://developers.yubico.com/yubikey-val/Validation_Protocol_V2.0.html
	sort.Strings(values)
	for _, v := range values {
		payload += v + "&"
	}
	payload = payload[:len(payload)-1]

	mac := hmac.New(sha1.New, key)
	mac.Write([]byte(payload))
	return mac.Sum(nil)
}

func loadKey(id string, dal *Dal) ([]byte, error) {
	i, err := dal.GetApp(id)
	if err != nil {
		return []byte{}, errors.New(NO_SUCH_CLIENT)
	}

	return i, nil
}

func reply(w http.ResponseWriter, otp, name, nonce, status, id string, dal *Dal) {
	values := []string{}
	key := []byte{}
	err := errors.New("")

	values = append(values, "nonce="+nonce)
	values = append(values, "otp="+otp)
	if status != MISSING_PARAMETER {
		key, err = loadKey(id, dal)
		if err == nil {
			values = append(values, "name="+name)
			values = append(values, "status="+status)
		} else {
			values = append(values, "status="+err.Error())
		}
	} else {
		values = append(values, "status="+status)
	}
	values = append(values, "t="+time.Now().Format(time.RFC3339))
	if status != MISSING_PARAMETER {
		values = append(values, "h="+base64.StdEncoding.EncodeToString(Sign(values, key)))
	}

	ret := ""
	for _, v := range values {
		ret += v + "\n"
	}

	w.Write([]byte(ret))
}

func runAPI(dal *Dal, host, port string) {
	r := mux.NewRouter()

	r.HandleFunc("/wsapi/2.0/verify", func(w http.ResponseWriter, r *http.Request) {
		checkOTP(w, r, dal)
	}).Methods("GET")

	http.Handle("/", r)
	log.Printf("Listening on: %s:%s...", host, port)
	http.ListenAndServe(host+":"+port, nil)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"fancyirc/types"

	"github.com/sorcix/irc"
)

var (
	servers = flag.String("servers",
		"localhost:8001",
		"(comma-separated) list of host:port network addresses of the server(s) to connect to")

	listen = flag.String("listen",
		"localhost:6667",
		"host:port to listen on for IRC connections")

	currentMaster string
	allServers    []string
)

const (
	pathCreateSession = "/fancyirc/v1/session"
	pathDeleteSession = "/fancyirc/v1/%s"
	pathPostMessage   = "/fancyirc/v1/%s/message"
	pathGetMessages   = "/fancyirc/v1/%s/messages?lastseen=%d"
)

// TODO(secure): persistent state:
// - the follow targets (channels)
// - the last known server(s) in the network. added to *servers
// - the last seen message id
// for hosted mode, this state is stored per-nickname, ideally encrypted with password

func sendFancyMessage(logPrefix, target, path string, data []byte) (*http.Response, error) {
	resp, err := http.Post(
		fmt.Sprintf("http://%s%s", target, path),
		"application/json",
		bytes.NewBuffer(data))

	if err != nil {
		// TODO(secure): try one of the other servers.
		return resp, err
	}

	if resp.StatusCode == http.StatusTemporaryRedirect {
		loc := resp.Header.Get("Location")
		if loc == "" {
			return resp, fmt.Errorf("Redirect has no Location header")
		}
		u, err := url.Parse(loc)
		if err != nil {
			return resp, fmt.Errorf("Could not parse redirection %q: %v", loc, err)
		}

		return sendFancyMessage(logPrefix, u.Host, path, data)
	}

	if resp.StatusCode != 200 {
		data, _ := ioutil.ReadAll(resp.Body)
		return resp, fmt.Errorf("sendFancyMessage(%s) failed with %v: %s", path, resp.Status, string(data))
	}

	log.Printf("%s ->fancy: %q\n", logPrefix, string(data))

	currentMaster = target
	return resp, nil
}

func sendIRCMessage(logPrefix string, ircConn *irc.Conn, msg irc.Message) {
	if err := ircConn.Encode(&msg); err != nil {
		log.Printf("%s Error sending IRC message %q: %v. Closing connection.\n", logPrefix, msg.Bytes(), err)
		// This leads to an error in .Decode(), terminating the handleIRC goroutine.
		ircConn.Close()
		return
	}
	log.Printf("%s ->irc: %q\n", logPrefix, msg.Bytes())
}

func createFancySession(logPrefix string) (session string, prefix irc.Prefix, err error) {
	var resp *http.Response
	resp, err = sendFancyMessage(logPrefix, currentMaster, pathCreateSession, []byte{})
	if err != nil {
		return
	}
	defer resp.Body.Close()

	type createSessionReply struct {
		Sessionid string
		Prefix    string
	}

	var createreply createSessionReply

	if err = json.NewDecoder(resp.Body).Decode(&createreply); err != nil {
		return
	}

	session = createreply.Sessionid
	prefix = irc.Prefix{Name: createreply.Prefix}
	return
}

func handleIRC(conn net.Conn) {
	var (
		logPrefix     = conn.RemoteAddr().String()
		ircConn       = irc.NewConn(conn)
		ircErrors     = make(chan error)
		ircMessages   = make(chan irc.Message)
		fancyMessages = make(chan string)

		ircPrefix irc.Prefix
		session   string
		quitmsg   string
		done      bool
		pingSent  bool
		err       error
	)

	session, ircPrefix, err = createFancySession(logPrefix)
	if err != nil {
		// TODO(secure): how can we properly kill an irc session? i.e. how do servers handle bans/overload?
		log.Printf("%s Could not create fancyirc session: %v\n", logPrefix, err)
		ircConn.Close()
		return
	}

	go func() {
		for {
			message, err := ircConn.Decode()
			if err != nil {
				ircErrors <- err
				return
			}
			log.Printf("%s <-irc: %q\n", logPrefix, message.Bytes())
			ircMessages <- *message
		}
	}()

	// TODO(secure): periodically get all the servers in the network, overwrite allServers (so that deletions work)

	go func() {
		var lastSeen types.FancyId

		for !done {
			host := allServers[rand.Intn(len(allServers))]
			// TODO(secure): exponential backoff in a per-server fashion
			log.Printf("%s Connecting to %q...\n", logPrefix, host)
			// TODO(secure): build targets (= filters) and add them to the url
			hostUrl := fmt.Sprintf("http://%s"+pathGetMessages, host, session, lastSeen)
			resp, err := http.Get(hostUrl)
			if err != nil {
				log.Printf("%s HTTP GET %q failed: %v\n", logPrefix, hostUrl, err)
				continue
			}

			if resp.StatusCode != 200 {
				log.Printf("%s Received unexpected status code from %q: %v\n", logPrefix, host, resp.Status)
				continue
			}

			// We set the host as currentMaster, not because the host is the
			// master, but because it is reachable. When sending messages, we will
			// either reach the master by chance or get redirected, at which point
			// we update currentMaster.
			currentMaster = host

			dec := json.NewDecoder(resp.Body)
			for !done {
				// TODO(secure): we need a ping message here as well, so that we can detect timeouts quickly. It could include the current servers.
				var msg types.FancyMessage
				if err := dec.Decode(&msg); err != nil {
					log.Printf("%s Protocol error on %q: Could not decode response chunk as JSON: %v\n", logPrefix, host, err)
					break
				}

				log.Printf("%s <-fancy: %q\n", logPrefix, msg.Data)
				fancyMessages <- msg.Data
				lastSeen = msg.Id
			}
			resp.Body.Close()
		}

		close(fancyMessages)

		log.Printf("Disconnecting.\n")
	}()

	// Read all remaining messages to prevent goroutine hangs.
	defer func() {
		for _ = range fancyMessages {
		}

		log.Printf("TODO: close session with msg %q\n", quitmsg)
		// TODO: close the session on the server
	}()

	for {
		select {
		case <-time.After(1 * time.Minute):
			// After no traffic in either direction for 1 minute, we send a PING
			// message. If a PING message was already sent, this means that we did
			// not receive a PONG message, so we close the connection with at
			// timeout.
			if pingSent {
				quitmsg = "ping timeout"
				ircConn.Close()
			} else {
				sendIRCMessage(logPrefix, ircConn, irc.Message{
					Prefix:  &ircPrefix,
					Command: irc.PING,
					Params:  []string{"fancyirc.proxy"},
				})
			}

		case err := <-ircErrors:
			log.Printf("Error in IRC client connection: %v\n", err)
			done = true
			return

		case msg := <-fancyMessages:
			if _, err := fmt.Fprintf(conn, "%s\n", msg); err != nil {
				log.Fatal(err)
			}

		case message := <-ircMessages:
			switch message.Command {
			case irc.PONG:
				log.Printf("%s received PONG reply.\n", logPrefix)
				pingSent = false
			case irc.PING:
				sendIRCMessage(logPrefix, ircConn, irc.Message{
					Prefix:  &ircPrefix,
					Command: irc.PONG,
					Params:  message.Params,
				})
			case irc.QUIT:
				quitmsg = message.Trailing
			default:
				//case irc.NICK:
				//	log.Printf("requested nickname is %q\n", message.Params[0])
				//	nick = message.Params[0]
				//	// TODO(secure): this needs to create a session on the IRC server.
				//	// TODO(secure): figure out whether we want to have at most 1 irc connection per nickname/session.
				//case irc.USER:
				//	// TODO(secure): the irc server, not the proxy, is supposed to send these messages
				//	// TODO(secure): send 002, 003, 004, 005, 251, 252, 254, 255, 265, 266, [motd = 375, 372, 376]
				//	reply(irc.Message{
				//		Command:  irc.RPL_WELCOME,
				//		Params:   []string{nick},
				//		Trailing: "Welcome to fancyirc :)",
				//	})
				//case irc.PRIVMSG:
				// TODO: we need to associate a session with this.
				if _, err := sendFancyMessage(logPrefix, currentMaster, fmt.Sprintf(pathPostMessage, session), message.Bytes()); err != nil {
					// TODO(secure): what should we do here?
					log.Printf("message could not be sent: %v\n", err)
				}
			}
		}
	}
}

func main() {
	flag.Parse()

	rand.Seed(time.Now().Unix())

	// Start with any server. Will be overwritten later.
	allServers = strings.Split(*servers, ",")
	if len(allServers) == 0 {
		log.Fatalf("Invalid -servers value (%q). Need at least one server.\n", *servers)
	}
	currentMaster = allServers[0]

	ln, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("fancyirc proxy listening on %q\n", *listen)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("Could not accept IRC client connection: %v\n", err)
			continue
		}
		go handleIRC(conn)
	}
}
<commit_msg>properly signal errors when creating a fancyirc session<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"fancyirc/types"

	"github.com/sorcix/irc"
)

var (
	servers = flag.String("servers",
		"localhost:8001",
		"(comma-separated) list of host:port network addresses of the server(s) to connect to")

	listen = flag.String("listen",
		"localhost:6667",
		"host:port to listen on for IRC connections")

	currentMaster string
	allServers    []string
)

const (
	pathCreateSession = "/fancyirc/v1/session"
	pathDeleteSession = "/fancyirc/v1/%s"
	pathPostMessage   = "/fancyirc/v1/%s/message"
	pathGetMessages   = "/fancyirc/v1/%s/messages?lastseen=%d"
)

// TODO(secure): persistent state:
// - the follow targets (channels)
// - the last known server(s) in the network. added to *servers
// - the last seen message id
// for hosted mode, this state is stored per-nickname, ideally encrypted with password

func sendFancyMessage(logPrefix, target, path string, data []byte) (*http.Response, error) {
	resp, err := http.Post(
		fmt.Sprintf("http://%s%s", target, path),
		"application/json",
		bytes.NewBuffer(data))

	if err != nil {
		// TODO(secure): try one of the other servers.
		return resp, err
	}

	if resp.StatusCode == http.StatusTemporaryRedirect {
		loc := resp.Header.Get("Location")
		if loc == "" {
			return resp, fmt.Errorf("Redirect has no Location header")
		}
		u, err := url.Parse(loc)
		if err != nil {
			return resp, fmt.Errorf("Could not parse redirection %q: %v", loc, err)
		}

		return sendFancyMessage(logPrefix, u.Host, path, data)
	}

	if resp.StatusCode != 200 {
		data, _ := ioutil.ReadAll(resp.Body)
		return resp, fmt.Errorf("sendFancyMessage(%s) failed with %v: %s", path, resp.Status, string(data))
	}

	log.Printf("%s ->fancy: %q\n", logPrefix, string(data))

	currentMaster = target
	return resp, nil
}

func sendIRCMessage(logPrefix string, ircConn *irc.Conn, msg irc.Message) {
	if err := ircConn.Encode(&msg); err != nil {
		log.Printf("%s Error sending IRC message %q: %v. Closing connection.\n", logPrefix, msg.Bytes(), err)
		// This leads to an error in .Decode(), terminating the handleIRC goroutine.
		ircConn.Close()
		return
	}
	log.Printf("%s ->irc: %q\n", logPrefix, msg.Bytes())
}

func createFancySession(logPrefix string) (session string, prefix irc.Prefix, err error) {
	var resp *http.Response
	resp, err = sendFancyMessage(logPrefix, currentMaster, pathCreateSession, []byte{})
	if err != nil {
		return
	}
	defer resp.Body.Close()

	type createSessionReply struct {
		Sessionid string
		Prefix    string
	}

	var createreply createSessionReply

	if err = json.NewDecoder(resp.Body).Decode(&createreply); err != nil {
		return
	}

	session = createreply.Sessionid
	prefix = irc.Prefix{Name: createreply.Prefix}
	return
}

func handleIRC(conn net.Conn) {
	var (
		logPrefix     = conn.RemoteAddr().String()
		ircConn       = irc.NewConn(conn)
		ircErrors     = make(chan error)
		ircMessages   = make(chan irc.Message)
		fancyMessages = make(chan string)

		ircPrefix irc.Prefix
		session   string
		quitmsg   string
		done      bool
		pingSent  bool
		err       error
	)

	session, ircPrefix, err = createFancySession(logPrefix)
	if err != nil {
		log.Printf("%s Could not create fancyirc session: %v\n", logPrefix, err)
		sendIRCMessage(logPrefix, ircConn, irc.Message{
			Command:  "ERROR",
			Trailing: fmt.Sprintf("Could not create fancyirc session: %v", err),
		})

		ircConn.Close()
		return
	}

	go func() {
		for {
			message, err := ircConn.Decode()
			if err != nil {
				ircErrors <- err
				return
			}
			log.Printf("%s <-irc: %q\n", logPrefix, message.Bytes())
			ircMessages <- *message
		}
	}()

	// TODO(secure): periodically get all the servers in the network, overwrite allServers (so that deletions work)

	go func() {
		var lastSeen types.FancyId

		for !done {
			host := allServers[rand.Intn(len(allServers))]
			// TODO(secure): exponential backoff in a per-server fashion
			log.Printf("%s Connecting to %q...\n", logPrefix, host)
			// TODO(secure): build targets (= filters) and add them to the url
			hostUrl := fmt.Sprintf("http://%s"+pathGetMessages, host, session, lastSeen)
			resp, err := http.Get(hostUrl)
			if err != nil {
				log.Printf("%s HTTP GET %q failed: %v\n", logPrefix, hostUrl, err)
				continue
			}

			if resp.StatusCode != 200 {
				log.Printf("%s Received unexpected status code from %q: %v\n", logPrefix, host, resp.Status)
				continue
			}

			// We set the host as currentMaster, not because the host is the
			// master, but because it is reachable. When sending messages, we will
			// either reach the master by chance or get redirected, at which point
			// we update currentMaster.
			currentMaster = host

			dec := json.NewDecoder(resp.Body)
			for !done {
				// TODO(secure): we need a ping message here as well, so that we can detect timeouts quickly. It could include the current servers.
				var msg types.FancyMessage
				if err := dec.Decode(&msg); err != nil {
					log.Printf("%s Protocol error on %q: Could not decode response chunk as JSON: %v\n", logPrefix, host, err)
					break
				}

				log.Printf("%s <-fancy: %q\n", logPrefix, msg.Data)
				fancyMessages <- msg.Data
				lastSeen = msg.Id
			}
			resp.Body.Close()
		}

		close(fancyMessages)

		log.Printf("Disconnecting.\n")
	}()

	// Read all remaining messages to prevent goroutine hangs.
	defer func() {
		for _ = range fancyMessages {
		}

		log.Printf("TODO: close session with msg %q\n", quitmsg)
		// TODO: close the session on the server
	}()

	for {
		select {
		case <-time.After(1 * time.Minute):
			// After no traffic in either direction for 1 minute, we send a PING
			// message. If a PING message was already sent, this means that we did
			// not receive a PONG message, so we close the connection with at
			// timeout.
			if pingSent {
				quitmsg = "ping timeout"
				ircConn.Close()
			} else {
				sendIRCMessage(logPrefix, ircConn, irc.Message{
					Prefix:  &ircPrefix,
					Command: irc.PING,
					Params:  []string{"fancyirc.proxy"},
				})
			}

		case err := <-ircErrors:
			log.Printf("Error in IRC client connection: %v\n", err)
			done = true
			return

		case msg := <-fancyMessages:
			if _, err := fmt.Fprintf(conn, "%s\n", msg); err != nil {
				log.Fatal(err)
			}

		case message := <-ircMessages:
			switch message.Command {
			case irc.PONG:
				log.Printf("%s received PONG reply.\n", logPrefix)
				pingSent = false
			case irc.PING:
				sendIRCMessage(logPrefix, ircConn, irc.Message{
					Prefix:  &ircPrefix,
					Command: irc.PONG,
					Params:  message.Params,
				})
			case irc.QUIT:
				quitmsg = message.Trailing
			default:
				//case irc.NICK:
				//	log.Printf("requested nickname is %q\n", message.Params[0])
				//	nick = message.Params[0]
				//	// TODO(secure): this needs to create a session on the IRC server.
				//	// TODO(secure): figure out whether we want to have at most 1 irc connection per nickname/session.
				//case irc.USER:
				//	// TODO(secure): the irc server, not the proxy, is supposed to send these messages
				//	// TODO(secure): send 002, 003, 004, 005, 251, 252, 254, 255, 265, 266, [motd = 375, 372, 376]
				//	reply(irc.Message{
				//		Command:  irc.RPL_WELCOME,
				//		Params:   []string{nick},
				//		Trailing: "Welcome to fancyirc :)",
				//	})
				//case irc.PRIVMSG:
				// TODO: we need to associate a session with this.
				if _, err := sendFancyMessage(logPrefix, currentMaster, fmt.Sprintf(pathPostMessage, session), message.Bytes()); err != nil {
					// TODO(secure): what should we do here?
					log.Printf("message could not be sent: %v\n", err)
				}
			}
		}
	}
}

func main() {
	flag.Parse()

	rand.Seed(time.Now().Unix())

	// Start with any server. Will be overwritten later.
	allServers = strings.Split(*servers, ",")
	if len(allServers) == 0 {
		log.Fatalf("Invalid -servers value (%q). Need at least one server.\n", *servers)
	}
	currentMaster = allServers[0]

	ln, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("fancyirc proxy listening on %q\n", *listen)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("Could not accept IRC client connection: %v\n", err)
			continue
		}
		go handleIRC(conn)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"fancyirc/types"

	"github.com/sorcix/irc"
)

var (
	servers = flag.String("servers",
		"localhost:8001",
		"(comma-separated) list of host:port network addresses of the server(s) to connect to")

	listen = flag.String("listen",
		"localhost:6667",
		"host:port to listen on for IRC connections")

	currentMaster string
	allServers    []string
)

const (
	pathCreateSession = "/fancyirc/v1/session"
	pathDeleteSession = "/fancyirc/v1/%s"
	pathPostMessage   = "/fancyirc/v1/%s/message"
	pathGetMessages   = "/fancyirc/v1/%s/messages?lastseen=%d"
)

// TODO(secure): persistent state:
// - the follow targets (channels)
// - the last known server(s) in the network. added to *servers
// - the last seen message id
// for hosted mode, this state is stored per-nickname, ideally encrypted with password

func sendFancyMessage(logPrefix, target, path string, data []byte) (*http.Response, error) {
	resp, err := http.Post(
		fmt.Sprintf("http://%s%s", target, path),
		"application/json",
		bytes.NewBuffer(data))

	if err != nil {
		// TODO(secure): try one of the other servers.
		return resp, err
	}

	if resp.StatusCode == http.StatusTemporaryRedirect {
		loc := resp.Header.Get("Location")
		if loc == "" {
			return resp, fmt.Errorf("Redirect has no Location header")
		}
		u, err := url.Parse(loc)
		if err != nil {
			return resp, fmt.Errorf("Could not parse redirection %q: %v", loc, err)
		}

		return sendFancyMessage(logPrefix, u.Host, path, data)
	}

	if resp.StatusCode != 200 {
		data, _ := ioutil.ReadAll(resp.Body)
		return resp, fmt.Errorf("sendFancyMessage(%s) failed with %v: %s", path, resp.Status, string(data))
	}

	log.Printf("%s ->fancy: %q\n", logPrefix, string(data))

	currentMaster = target
	return resp, nil
}

func sendIRCMessage(logPrefix string, ircConn *irc.Conn, msg irc.Message) {
	if err := ircConn.Encode(&msg); err != nil {
		log.Printf("%s Error sending IRC message %q: %v. Closing connection.\n", logPrefix, msg.Bytes(), err)
		// This leads to an error in .Decode(), terminating the handleIRC goroutine.
		ircConn.Close()
		return
	}
	log.Printf("%s ->irc: %q\n", logPrefix, msg.Bytes())
}

func createFancySession(logPrefix string) (session string, prefix irc.Prefix, err error) {
	var resp *http.Response
	resp, err = sendFancyMessage(logPrefix, currentMaster, pathCreateSession, []byte{})
	if err != nil {
		return
	}
	defer resp.Body.Close()

	type createSessionReply struct {
		Sessionid string
		Prefix    string
	}

	var createreply createSessionReply

	if err = json.NewDecoder(resp.Body).Decode(&createreply); err != nil {
		return
	}

	session = createreply.Sessionid
	prefix = irc.Prefix{Name: createreply.Prefix}
	return
}

func handleIRC(conn net.Conn) {
	var (
		logPrefix     = conn.RemoteAddr().String()
		ircConn       = irc.NewConn(conn)
		ircErrors     = make(chan error)
		ircMessages   = make(chan irc.Message)
		fancyMessages = make(chan string)

		ircPrefix irc.Prefix
		session   string
		quitmsg   string
		done      bool
		pingSent  bool
		err       error
	)

	session, ircPrefix, err = createFancySession(logPrefix)
	if err != nil {
		log.Printf("%s Could not create fancyirc session: %v\n", logPrefix, err)
		sendIRCMessage(logPrefix, ircConn, irc.Message{
			Command:  "ERROR",
			Trailing: fmt.Sprintf("Could not create fancyirc session: %v", err),
		})

		ircConn.Close()
		return
	}

	go func() {
		for {
			message, err := ircConn.Decode()
			if err != nil {
				ircErrors <- err
				return
			}
			log.Printf("%s <-irc: %q\n", logPrefix, message.Bytes())
			ircMessages <- *message
		}
	}()

	// TODO(secure): periodically get all the servers in the network, overwrite allServers (so that deletions work)

	go func() {
		var lastSeen types.FancyId

		for !done {
			host := allServers[rand.Intn(len(allServers))]
			// TODO(secure): exponential backoff in a per-server fashion
			log.Printf("%s Connecting to %q...\n", logPrefix, host)
			// TODO(secure): build targets (= filters) and add them to the url
			hostUrl := fmt.Sprintf("http://%s"+pathGetMessages, host, session, lastSeen)
			resp, err := http.Get(hostUrl)
			if err != nil {
				log.Printf("%s HTTP GET %q failed: %v\n", logPrefix, hostUrl, err)
				continue
			}

			if resp.StatusCode != 200 {
				log.Printf("%s Received unexpected status code from %q: %v\n", logPrefix, host, resp.Status)
				continue
			}

			// We set the host as currentMaster, not because the host is the
			// master, but because it is reachable. When sending messages, we will
			// either reach the master by chance or get redirected, at which point
			// we update currentMaster.
			currentMaster = host

			dec := json.NewDecoder(resp.Body)
			for !done {
				// TODO(secure): we need a ping message here as well, so that we can detect timeouts quickly. It could include the current servers.
				var msg types.FancyMessage
				if err := dec.Decode(&msg); err != nil {
					log.Printf("%s Protocol error on %q: Could not decode response chunk as JSON: %v\n", logPrefix, host, err)
					break
				}

				log.Printf("%s <-fancy: %q\n", logPrefix, msg.Data)
				fancyMessages <- msg.Data
				lastSeen = msg.Id
			}
			resp.Body.Close()
		}

		close(fancyMessages)

		log.Printf("Disconnecting.\n")
	}()

	// Read all remaining messages to prevent goroutine hangs.
	defer func() {
		for _ = range fancyMessages {
		}

		log.Printf("TODO: close session with msg %q\n", quitmsg)
		// TODO: close the session on the server
	}()

	for {
		select {
		case <-time.After(1 * time.Minute):
			// After no traffic in either direction for 1 minute, we send a PING
			// message. If a PING message was already sent, this means that we did
			// not receive a PONG message, so we close the connection with at
			// timeout.
			if pingSent {
				quitmsg = "ping timeout"
				ircConn.Close()
			} else {
				sendIRCMessage(logPrefix, ircConn, irc.Message{
					Prefix:  &ircPrefix,
					Command: irc.PING,
					Params:  []string{"fancyirc.proxy"},
				})
			}

		case err := <-ircErrors:
			log.Printf("Error in IRC client connection: %v\n", err)
			done = true
			return

		case msg := <-fancyMessages:
			if _, err := fmt.Fprintf(conn, "%s\n", msg); err != nil {
				log.Fatal(err)
			}

		case message := <-ircMessages:
			switch message.Command {
			case irc.PONG:
				log.Printf("%s received PONG reply.\n", logPrefix)
				pingSent = false
			case irc.PING:
				sendIRCMessage(logPrefix, ircConn, irc.Message{
					Prefix:  &ircPrefix,
					Command: irc.PONG,
					Params:  message.Params,
				})
			case irc.QUIT:
				quitmsg = message.Trailing
			default:
				//case irc.NICK:
				//	log.Printf("requested nickname is %q\n", message.Params[0])
				//	nick = message.Params[0]
				//	// TODO(secure): this needs to create a session on the IRC server.
				//	// TODO(secure): figure out whether we want to have at most 1 irc connection per nickname/session.
				//case irc.USER:
				//	// TODO(secure): the irc server, not the proxy, is supposed to send these messages
				//	// TODO(secure): send 002, 003, 004, 005, 251, 252, 254, 255, 265, 266, [motd = 375, 372, 376]
				//	reply(irc.Message{
				//		Command:  irc.RPL_WELCOME,
				//		Params:   []string{nick},
				//		Trailing: "Welcome to fancyirc :)",
				//	})
				//case irc.PRIVMSG:
				// TODO: we need to associate a session with this.
				if _, err := sendFancyMessage(logPrefix, currentMaster, fmt.Sprintf(pathPostMessage, session), message.Bytes()); err != nil {
					// TODO(secure): what should we do here?
					log.Printf("message could not be sent: %v\n", err)
				}
			}
		}
	}
}

func main() {
	flag.Parse()

	rand.Seed(time.Now().Unix())

	// Start with any server. Will be overwritten later.
	allServers = strings.Split(*servers, ",")
	if len(allServers) == 0 {
		log.Fatalf("Invalid -servers value (%q). Need at least one server.\n", *servers)
	}
	currentMaster = allServers[0]

	ln, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("fancyirc proxy listening on %q\n", *listen)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("Could not accept IRC client connection: %v\n", err)
			continue
		}
		go handleIRC(conn)
	}
}
<commit_msg>update comments<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"fancyirc/types"

	"github.com/sorcix/irc"
)

var (
	servers = flag.String("servers",
		"localhost:8001",
		"(comma-separated) list of host:port network addresses of the server(s) to connect to")

	listen = flag.String("listen",
		"localhost:6667",
		"host:port to listen on for IRC connections")

	currentMaster string
	allServers    []string
)

const (
	pathCreateSession = "/fancyirc/v1/session"
	pathDeleteSession = "/fancyirc/v1/%s"
	pathPostMessage   = "/fancyirc/v1/%s/message"
	pathGetMessages   = "/fancyirc/v1/%s/messages?lastseen=%d"
)

// TODO(secure): persistent state:
// - the last known server(s) in the network. added to *servers
// - for resuming sessions (later): the last seen message id, perhaps setup messages (JOINs, MODEs, …)
// for hosted mode, this state is stored per-nickname, ideally encrypted with password

func sendFancyMessage(logPrefix, target, path string, data []byte) (*http.Response, error) {
	resp, err := http.Post(
		fmt.Sprintf("http://%s%s", target, path),
		"application/json",
		bytes.NewBuffer(data))

	if err != nil {
		// TODO(secure): try one of the other servers.
		return resp, err
	}

	if resp.StatusCode == http.StatusTemporaryRedirect {
		loc := resp.Header.Get("Location")
		if loc == "" {
			return resp, fmt.Errorf("Redirect has no Location header")
		}
		u, err := url.Parse(loc)
		if err != nil {
			return resp, fmt.Errorf("Could not parse redirection %q: %v", loc, err)
		}

		return sendFancyMessage(logPrefix, u.Host, path, data)
	}

	if resp.StatusCode != 200 {
		data, _ := ioutil.ReadAll(resp.Body)
		return resp, fmt.Errorf("sendFancyMessage(%s) failed with %v: %s", path, resp.Status, string(data))
	}

	log.Printf("%s ->fancy: %q\n", logPrefix, string(data))

	currentMaster = target
	return resp, nil
}

func sendIRCMessage(logPrefix string, ircConn *irc.Conn, msg irc.Message) {
	if err := ircConn.Encode(&msg); err != nil {
		log.Printf("%s Error sending IRC message %q: %v. Closing connection.\n", logPrefix, msg.Bytes(), err)
		// This leads to an error in .Decode(), terminating the handleIRC goroutine.
		ircConn.Close()
		return
	}
	log.Printf("%s ->irc: %q\n", logPrefix, msg.Bytes())
}

func createFancySession(logPrefix string) (session string, prefix irc.Prefix, err error) {
	var resp *http.Response
	resp, err = sendFancyMessage(logPrefix, currentMaster, pathCreateSession, []byte{})
	if err != nil {
		return
	}
	defer resp.Body.Close()

	type createSessionReply struct {
		Sessionid string
		Prefix    string
	}

	var createreply createSessionReply

	if err = json.NewDecoder(resp.Body).Decode(&createreply); err != nil {
		return
	}

	session = createreply.Sessionid
	prefix = irc.Prefix{Name: createreply.Prefix}
	return
}

func handleIRC(conn net.Conn) {
	var (
		logPrefix     = conn.RemoteAddr().String()
		ircConn       = irc.NewConn(conn)
		ircErrors     = make(chan error)
		ircMessages   = make(chan irc.Message)
		fancyMessages = make(chan string)

		ircPrefix irc.Prefix
		session   string
		quitmsg   string
		done      bool
		pingSent  bool
		err       error
	)

	session, ircPrefix, err = createFancySession(logPrefix)
	if err != nil {
		log.Printf("%s Could not create fancyirc session: %v\n", logPrefix, err)
		sendIRCMessage(logPrefix, ircConn, irc.Message{
			Command:  "ERROR",
			Trailing: fmt.Sprintf("Could not create fancyirc session: %v", err),
		})

		ircConn.Close()
		return
	}

	go func() {
		for {
			message, err := ircConn.Decode()
			if err != nil {
				ircErrors <- err
				return
			}
			log.Printf("%s <-irc: %q\n", logPrefix, message.Bytes())
			ircMessages <- *message
		}
	}()

	// TODO(secure): periodically get all the servers in the network, overwrite allServers (so that deletions work)

	go func() {
		var lastSeen types.FancyId

		for !done {
			host := allServers[rand.Intn(len(allServers))]
			// TODO(secure): exponential backoff in a per-server fashion
			log.Printf("%s Connecting to %q...\n", logPrefix, host)
			hostUrl := fmt.Sprintf("http://%s"+pathGetMessages, host, session, lastSeen)
			resp, err := http.Get(hostUrl)
			if err != nil {
				log.Printf("%s HTTP GET %q failed: %v\n", logPrefix, hostUrl, err)
				continue
			}

			if resp.StatusCode != 200 {
				log.Printf("%s Received unexpected status code from %q: %v\n", logPrefix, host, resp.Status)
				continue
			}

			// We set the host as currentMaster, not because the host is the
			// master, but because it is reachable. When sending messages, we will
			// either reach the master by chance or get redirected, at which point
			// we update currentMaster.
			currentMaster = host

			dec := json.NewDecoder(resp.Body)
			for !done {
				// TODO(secure): we need a ping message here as well, so that we can detect timeouts quickly. It could include the current servers.
				var msg types.FancyMessage
				if err := dec.Decode(&msg); err != nil {
					log.Printf("%s Protocol error on %q: Could not decode response chunk as JSON: %v\n", logPrefix, host, err)
					break
				}

				log.Printf("%s <-fancy: %q\n", logPrefix, msg.Data)
				fancyMessages <- msg.Data
				lastSeen = msg.Id
			}
			resp.Body.Close()
		}

		close(fancyMessages)

		log.Printf("Disconnecting.\n")
	}()

	// Read all remaining messages to prevent goroutine hangs.
	defer func() {
		for _ = range fancyMessages {
		}

		log.Printf("TODO: close session with msg %q\n", quitmsg)
		// TODO: close the session on the server
	}()

	for {
		select {
		case <-time.After(1 * time.Minute):
			// After no traffic in either direction for 1 minute, we send a PING
			// message. If a PING message was already sent, this means that we did
			// not receive a PONG message, so we close the connection with at
			// timeout.
			if pingSent {
				quitmsg = "ping timeout"
				ircConn.Close()
			} else {
				sendIRCMessage(logPrefix, ircConn, irc.Message{
					Prefix:  &ircPrefix,
					Command: irc.PING,
					Params:  []string{"fancyirc.proxy"},
				})
			}

		case err := <-ircErrors:
			log.Printf("Error in IRC client connection: %v\n", err)
			done = true
			return

		case msg := <-fancyMessages:
			if _, err := fmt.Fprintf(conn, "%s\n", msg); err != nil {
				log.Fatal(err)
			}

		case message := <-ircMessages:
			switch message.Command {
			case irc.PONG:
				log.Printf("%s received PONG reply.\n", logPrefix)
				pingSent = false
			case irc.PING:
				sendIRCMessage(logPrefix, ircConn, irc.Message{
					Prefix:  &ircPrefix,
					Command: irc.PONG,
					Params:  message.Params,
				})
			case irc.QUIT:
				quitmsg = message.Trailing
			default:
				if _, err := sendFancyMessage(logPrefix, currentMaster, fmt.Sprintf(pathPostMessage, session), message.Bytes()); err != nil {
					// TODO(secure): what should we do here?
					log.Printf("message could not be sent: %v\n", err)
				}
			}
		}
	}
}

func main() {
	flag.Parse()

	rand.Seed(time.Now().Unix())

	// Start with any server. Will be overwritten later.
	allServers = strings.Split(*servers, ",")
	if len(allServers) == 0 {
		log.Fatalf("Invalid -servers value (%q). Need at least one server.\n", *servers)
	}
	currentMaster = allServers[0]

	ln, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("fancyirc proxy listening on %q\n", *listen)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("Could not accept IRC client connection: %v\n", err)
			continue
		}
		go handleIRC(conn)
	}
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
		}
		go p.start(powerCallback)
	}
}

func getResolvedAddresses(localHost, remoteHost *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", *remoteHost)
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	if islocal {
		for {
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			fmt.Printf("%#v\n", buff[:n])
			if remainingBytes > 0 {
				if remainingBytes <= n {
					newPacket = true
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = n - remainingBytes
					fmt.Println(msg)
				} else {
					newPacket = false
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = remainingBytes - n
				}
			}
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
		NewP:
			if newPacket {
				remainingBytes = 0
				newPacket = false
				msg = ""
				t := r.byte()
				n = n - 1
				fmt.Println(t)
				switch t {
				case query:
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					remainingBytes = remainingBytes - 4
					if remainingBytes > 0 {
						if remainingBytes <= n {
							newPacket = true
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = n - remainingBytes
							fmt.Println(msg)
							goto NewP
						} else {
							newPacket = false
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = remainingBytes - n
						}
						fmt.Printf("Remaining bytes: %d\n", remainingBytes)
					}
				// case rowDescription:
				// case dataRow:
				// case bindComplete:
				// case commandComplete:
				// 	commandTag = CommandTag(r.readCString())
				default:
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			}
			r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// fmt.Printf("t: %#v\n", t)

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}
<commit_msg>Update<commit_after>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
		}
		go p.start(powerCallback)
	}
}

func getResolvedAddresses(localHost, remoteHost *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", *remoteHost)
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	// newPacket := true
	// var msg string
	// remainingBytes := 0
	if islocal {
		for {
			// var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			// 	r = buff[:n]
			// 	fmt.Printf("%#v\n", buff[:n])
			// 	if remainingBytes > 0 {
			// 		if remainingBytes <= n {
			// 			newPacket = true
			// 			msg = msg + string(r.next(remainingBytes))
			// 			remainingBytes = n - remainingBytes
			// 			fmt.Println(msg)
			// 		} else {
			// 			newPacket = false
			// 			msg = msg + string(r.next(remainingBytes))
			// 			remainingBytes = remainingBytes - n
			// 		}
			// 	}
			// 	fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			// NewP:
			// 	if newPacket {
			// 		remainingBytes = 0
			// 		newPacket = false
			// 		msg = ""
			// 		t := r.byte()
			// 		n = n - 1
			// 		fmt.Println(t)
			// 		switch t {
			// 		case query:
			// 			// c.rxReadyForQuery(r)
			// 			remainingBytes = r.int32()
			// 			remainingBytes = remainingBytes - 4
			// 			if remainingBytes > 0 {
			// 				if remainingBytes <= n {
			// 					newPacket = true
			// 					msg = msg + string(r.next(remainingBytes))
			// 					remainingBytes = n - remainingBytes
			// 					fmt.Println(msg)
			// 					goto NewP
			// 				} else {
			// 					newPacket = false
			// 					msg = msg + string(r.next(remainingBytes))
			// 					remainingBytes = remainingBytes - n
			// 				}
			// 				fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			// 			}
			// 		// case rowDescription:
			// 		// case dataRow:
			// 		// case bindComplete:
			// 		// case commandComplete:
			// 		// 	commandTag = CommandTag(r.readCString())
			// 		default:
			// 			// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
			// 			// 	softErr = e
			// 			// }
			// 		}
			// 	}
			// 	r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// fmt.Printf("t: %#v\n", t)

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}
<|endoftext|>


<commit_before>package proxy

import (
	disc "github.com/jeffjen/go-discovery"
	"github.com/jeffjen/go-libkv/libkv"
	"github.com/jeffjen/go-proxy/proxy"

	log "github.com/Sirupsen/logrus"
	etcd "github.com/coreos/etcd/client"
	ctx "golang.org/x/net/context"

	"encoding/json"
)

var (
	ProxyConfigKey string

	ConfigReset ctx.CancelFunc

	retry = &proxy.Backoff{}
)

func parse(spec string) (*Info, error) {
	var i = new(Info)
	if err := json.Unmarshal([]byte(spec), i); err != nil {
		return nil, err
	}
	if i.Name == "" {
		return nil, ErrMissingName
	}
	return i, nil
}

func get(value string) (targets []*Info) {
	targets = make([]*Info, 0)
	if err := json.Unmarshal([]byte(value), &targets); err != nil {
		log.WithFields(log.Fields{"err": err}).Warning("bad proxy spec")
		targets = nil
	}
	return
}

func doReload(pxycfg []*Info) {
	it, mod := Store.IterateW()
	for elem := range it {
		mod <- &libkv.Value{R: true}
		elem.X.(*Info).Cancel()
	}
	for _, spec := range Targets {
		meta, err := parse(spec)
		if err != nil {
			log.WithFields(log.Fields{"err": err}).Warning("reload")
			continue
		}
		if err = Listen(meta); err != nil {
			if err != ErrProxyExist {
				log.WithFields(log.Fields{"err": err}).Debug("reload")
			}
		}
	}
	for _, meta := range pxycfg {
		if err := Listen(meta); err != nil {
			if err != ErrProxyExist {
				log.WithFields(log.Fields{"err": err}).Debug("reload")
			}
		}
	}
}

func reloadWorker() chan<- []*Info {
	order := make(chan []*Info)
	go func() {
		for o := range order {
			doReload(o)
		}
	}()
	return order
}

func doWatch(c ctx.Context, watcher etcd.Watcher) <-chan []*Info {
	v := make(chan []*Info)
	go func() {
		evt, err := watcher.Next(c)
		if err != nil {
			log.WithFields(log.Fields{"err": err}).Debug("config")
			retry.Delay()
			v <- nil
		} else {
			retry.Reset()
			if evt.Node.Dir {
				log.WithFields(log.Fields{"key": evt.Node.Key}).Warning("not a valid node")
				v <- nil
			} else {
				switch evt.Action {
				default:
					v <- nil
					break
				case "set":
					v <- get(evt.Node.Value)
					break
				case "del":
				case "expire":
					v <- make([]*Info, 0)
					break
				}
			}
		}
	}()
	return v
}

func followBootStrap() {
	cfg := etcd.Config{Endpoints: disc.Endpoints()}
	kAPI, err := proxy.NewKeysAPI(cfg)
	if err != nil {
		log.WithFields(log.Fields{"err": err}).Warning("bootstrap")
		return
	}
	resp, err := kAPI.Get(RootContext, ProxyConfigKey, nil)
	if err != nil {
		log.WithFields(log.Fields{"err": err}).Warning("bootstrap")
		doReload(make([]*Info, 0))
	} else if resp.Node.Dir {
		log.WithFields(log.Fields{"key": resp.Node.Key}).Warning("not a valid node")
		doReload(make([]*Info, 0))
	} else {
		log.WithFields(log.Fields{"key": resp.Node.Key, "val": resp.Node.Value}).Debug("cfgkey")
		if pxycfg := get(resp.Node.Value); pxycfg != nil {
			doReload(pxycfg)
		} else {
			doReload(make([]*Info, 0))
		}
	}
}

func Follow() {
	followBootStrap() // bootstrap proxy config

	var c ctx.Context

	c, ConfigReset = ctx.WithCancel(RootContext)
	go func() {
		cfg := etcd.Config{Endpoints: disc.Endpoints()}
		watcher, err := proxy.NewWatcher(cfg, ProxyConfigKey, 0)
		if err != nil {
			log.WithFields(log.Fields{"err": err}).Warning("config")
			return
		}
		order := reloadWorker()
		defer close(order)
		for yay := true; yay; {
			v := doWatch(c, watcher)
			select {
			case <-c.Done():
				yay = false
			case proxyTargets, ok := <-v:
				if ok && proxyTargets != nil {
					order <- proxyTargets
				}
				yay = ok
			}
		}
	}()
}
<commit_msg>NIT: log reload<commit_after>package proxy

import (
	disc "github.com/jeffjen/go-discovery"
	"github.com/jeffjen/go-libkv/libkv"
	"github.com/jeffjen/go-proxy/proxy"

	log "github.com/Sirupsen/logrus"
	etcd "github.com/coreos/etcd/client"
	ctx "golang.org/x/net/context"

	"encoding/json"
)

var (
	ProxyConfigKey string

	ConfigReset ctx.CancelFunc

	retry = &proxy.Backoff{}
)

func parse(spec string) (*Info, error) {
	var i = new(Info)
	if err := json.Unmarshal([]byte(spec), i); err != nil {
		return nil, err
	}
	if i.Name == "" {
		return nil, ErrMissingName
	}
	return i, nil
}

func get(value string) (targets []*Info) {
	targets = make([]*Info, 0)
	if err := json.Unmarshal([]byte(value), &targets); err != nil {
		log.WithFields(log.Fields{"err": err}).Warning("bad proxy spec")
		targets = nil
	}
	return
}

func doReload(pxycfg []*Info) {
	it, mod := Store.IterateW()
	for elem := range it {
		mod <- &libkv.Value{R: true}
		elem.X.(*Info).Cancel()
	}
	log.WithFields(log.Fields{"count": len(Targets)}).Debug("reload from args")
	for _, spec := range Targets {
		meta, err := parse(spec)
		if err != nil {
			log.WithFields(log.Fields{"err": err}).Warning("reload")
			continue
		}
		if err = Listen(meta); err != nil {
			if err != ErrProxyExist {
				log.WithFields(log.Fields{"err": err}).Debug("reload")
			}
		}
	}
	log.WithFields(log.Fields{"count": len(pxycfg)}).Debug("reload from cfgkey")
	for _, meta := range pxycfg {
		if err := Listen(meta); err != nil {
			if err != ErrProxyExist {
				log.WithFields(log.Fields{"err": err}).Debug("reload")
			}
		}
	}
}

func reloadWorker() chan<- []*Info {
	order := make(chan []*Info)
	go func() {
		for o := range order {
			doReload(o)
		}
	}()
	return order
}

func doWatch(c ctx.Context, watcher etcd.Watcher) <-chan []*Info {
	v := make(chan []*Info)
	go func() {
		evt, err := watcher.Next(c)
		if err != nil {
			log.WithFields(log.Fields{"err": err}).Debug("config")
			retry.Delay()
			v <- nil
		} else {
			retry.Reset()
			if evt.Node.Dir {
				log.WithFields(log.Fields{"key": evt.Node.Key}).Warning("not a valid node")
				v <- nil
			} else {
				switch evt.Action {
				default:
					v <- nil
					break
				case "set":
					v <- get(evt.Node.Value)
					break
				case "del":
				case "expire":
					v <- make([]*Info, 0)
					break
				}
			}
		}
	}()
	return v
}

func followBootStrap() {
	cfg := etcd.Config{Endpoints: disc.Endpoints()}
	kAPI, err := proxy.NewKeysAPI(cfg)
	if err != nil {
		log.WithFields(log.Fields{"err": err}).Warning("bootstrap")
		return
	}
	resp, err := kAPI.Get(RootContext, ProxyConfigKey, nil)
	if err != nil {
		log.WithFields(log.Fields{"err": err}).Warning("bootstrap")
		doReload(make([]*Info, 0))
	} else if resp.Node.Dir {
		log.WithFields(log.Fields{"key": resp.Node.Key}).Warning("not a valid node")
		doReload(make([]*Info, 0))
	} else {
		log.WithFields(log.Fields{"key": resp.Node.Key, "val": resp.Node.Value}).Debug("cfgkey")
		if pxycfg := get(resp.Node.Value); pxycfg != nil {
			doReload(pxycfg)
		} else {
			doReload(make([]*Info, 0))
		}
	}
}

func Follow() {
	followBootStrap() // bootstrap proxy config

	var c ctx.Context

	c, ConfigReset = ctx.WithCancel(RootContext)
	go func() {
		cfg := etcd.Config{Endpoints: disc.Endpoints()}
		watcher, err := proxy.NewWatcher(cfg, ProxyConfigKey, 0)
		if err != nil {
			log.WithFields(log.Fields{"err": err}).Warning("config")
			return
		}
		order := reloadWorker()
		defer close(order)
		for yay := true; yay; {
			v := doWatch(c, watcher)
			select {
			case <-c.Done():
				yay = false
			case proxyTargets, ok := <-v:
				if ok && proxyTargets != nil {
					order <- proxyTargets
				}
				yay = ok
			}
		}
	}()
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package psbt

// signer encapsulates the role 'Signer'
// as specified in BIP174; it controls the insertion of signatures;
// the Sign() function will attempt to insert signatures using
// Updater.addPartialSignature, after first ensuring the Psbt is in the
// correct state.

import (
	"github.com/btcsuite/btcd/txscript"
)

// Sign allows the caller to sign a PSBT at a particular input; they
// must provide a signature and a pubkey, both as byte slices; they can also
// optionally provide both witnessScript and/or redeemScript, otherwise
// these arguments must be set as nil (and in that case, they must already
// be present in the PSBT if required for signing to succeed).
//
// Return value:
// 0 indicates that the partial signature was successfully attached.
// 1 indicates that this input is already finalized, so the provided
// signature was *not* attached
// -1 indicates that the provided signature data was not valid. In this
// case an error will also be returned.
//
// This serves as a wrapper around Updater.addPartialSignature;
// it ensures that the redeemScript and witnessScript are updated as needed
// (note that the Updater is allowed to add redeemScripts and witnessScripts
// independently, before signing), and ensures that the right form of utxo
// field (NonWitnessUtxo or WitnessUtxo) is included in the input so that
// signature insertion (and then finalization) can take place.
func (u *Updater) Sign(inIndex int, sig []byte, pubKey []byte,
	redeemScript []byte, witnessScript []byte) (int, error) {

	if isFinalized(u.Upsbt, inIndex) {
		return 1, nil
	}

	if witnessScript != nil {
		// Add the witnessScript to the PSBT in preparation.
		// If it already exists, it will be overwritten.
		err := u.AddInWitnessScript(witnessScript, inIndex)
		if err != nil {
			return -1, err
		}
	}

	if redeemScript != nil {
		// Add the redeemScript to the PSBT in preparation.
		// If it already exists, it will be overwritten.
		err := u.AddInRedeemScript(redeemScript, inIndex)
		if err != nil {
			return -1, err
		}
	}

	// At this point, the PSBT must have the requisite
	// witnessScript or redeemScript fields for signing to succeed.

	// case 1: if witnessScript is present, it must be of type witness;
	// if not, signature insertion will of course fail.
	if u.Upsbt.Inputs[inIndex].WitnessScript != nil {
		if u.Upsbt.Inputs[inIndex].WitnessUtxo == nil {
			err := nonWitnessToWitness(u.Upsbt, inIndex)
			if err != nil {
				return -1, err
			}
		}
		err := u.addPartialSignature(inIndex, sig, pubKey)
		if err != nil {
			return -1, err
		}
	} else if u.Upsbt.Inputs[inIndex].RedeemScript != nil {
		// case 2: no witness script, only redeem script; can be legacy
		// p2sh or p2sh-wrapped p2wkh
		// We only need to decide if the input is witness, and we don't
		// rely on the witnessutxo/nonwitnessutxo in the PSBT, instead
		// we check the redeemScript content:
		if txscript.IsWitnessProgram(redeemScript) {
			if u.Upsbt.Inputs[inIndex].WitnessUtxo == nil {
				err := nonWitnessToWitness(u.Upsbt, inIndex)
				if err != nil {
					return -1, err
				}
			}
		}
		// If it is not a valid witness program, we here assume
		// that the provided WitnessUtxo/NonWitnessUtxo field was correct.
		err := u.addPartialSignature(inIndex, sig, pubKey)
		if err != nil {
			return -1, err
		}
	} else {
		// case 3: Neither provided only works for native p2wkh, or
		// non-segwit non-p2sh. To check if it's segwit, check
		// the scriptPubKey of the output.
		if u.Upsbt.Inputs[inIndex].WitnessUtxo == nil {
			outIndex := u.Upsbt.UnsignedTx.TxIn[inIndex].
				PreviousOutPoint.Index
			script := u.Upsbt.Inputs[inIndex].NonWitnessUtxo.
				TxOut[outIndex].PkScript
			if txscript.IsWitnessProgram(script) {
				err := nonWitnessToWitness(u.Upsbt, inIndex)
				if err != nil {
					return -1, err
				}
			}
		}
		err := u.addPartialSignature(inIndex, sig, pubKey)
		if err != nil {
			return -1, err
		}
	}
	return 0, nil
}

// nonWitnessToWitness extracts the TxOut from the existing
// NonWitnessUtxo field in the given PSBT input and sets it as type
// witness by replacing the NonWitnessUtxo field with a WitnessUtxo
// field. See https://github.com/bitcoin/bitcoin/pull/14197
func nonWitnessToWitness(p *Psbt, inIndex int) error {
	outIndex := p.UnsignedTx.TxIn[inIndex].PreviousOutPoint.Index
	txout := p.Inputs[inIndex].NonWitnessUtxo.TxOut[outIndex]
	// Remove the non-witness first, else sanity check will not pass:
	p.Inputs[inIndex].NonWitnessUtxo = nil
	u := Updater{Upsbt: p}
	return u.AddInWitnessUtxo(txout, inIndex)
}
<commit_msg>psbt: create new enum type for return values of Sign method<commit_after>// Copyright (c) 2018 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package psbt

// signer encapsulates the role 'Signer' as specified in BIP174; it controls
// the insertion of signatures; the Sign() function will attempt to insert
// signatures using Updater.addPartialSignature, after first ensuring the Psbt
// is in the correct state.

import (
	"github.com/btcsuite/btcd/txscript"
)

// SignOutcome is a enum-like value that expresses the outcome of a call to the
// Sign method.
type SignOutcome int

const (
	// SignSuccesful indicates that the partial signature was successfully
	// attached.
	SignSuccesful = 0

	// SignFinalized  indicates that this input is already finalized, so the provided
	// signature was *not* attached
	SignFinalized = 1

	// SignInvalid indicates that the provided signature data was not valid. In this case
	// an error will also be returned.
	SignInvalid = -1
)

// Sign allows the caller to sign a PSBT at a particular input; they
// must provide a signature and a pubkey, both as byte slices; they can also
// optionally provide both witnessScript and/or redeemScript, otherwise these
// arguments must be set as nil (and in that case, they must already be present
// in the PSBT if required for signing to succeed).
//
// This serves as a wrapper around Updater.addPartialSignature; it ensures that
// the redeemScript and witnessScript are updated as needed (note that the
// Updater is allowed to add redeemScripts and witnessScripts independently,
// before signing), and ensures that the right form of utxo field
// (NonWitnessUtxo or WitnessUtxo) is included in the input so that signature
// insertion (and then finalization) can take place.
func (u *Updater) Sign(inIndex int, sig []byte, pubKey []byte,
	redeemScript []byte, witnessScript []byte) (SignOutcome, error) {

	if isFinalized(u.Upsbt, inIndex) {
		return SignFinalized, nil
	}

	// Add the witnessScript to the PSBT in preparation.  If it already
	// exists, it will be overwritten.
	if witnessScript != nil {
		err := u.AddInWitnessScript(witnessScript, inIndex)
		if err != nil {
			return SignInvalid, err
		}
	}

	// Add the redeemScript to the PSBT in preparation.  If it already
	// exists, it will be overwritten.
	if redeemScript != nil {
		err := u.AddInRedeemScript(redeemScript, inIndex)
		if err != nil {
			return SignInvalid, err
		}
	}

	// At this point, the PSBT must have the requisite witnessScript or
	// redeemScript fields for signing to succeed.
	//
	// Case 1: if witnessScript is present, it must be of type witness;
	// if not, signature insertion will of course fail.
	switch {
	case u.Upsbt.Inputs[inIndex].WitnessScript != nil:
		if u.Upsbt.Inputs[inIndex].WitnessUtxo == nil {
			err := nonWitnessToWitness(u.Upsbt, inIndex)
			if err != nil {
				return SignInvalid, err
			}
		}

		err := u.addPartialSignature(inIndex, sig, pubKey)
		if err != nil {
			return SignInvalid, err
		}

	// Case 2: no witness script, only redeem script; can be legacy p2sh or
	// p2sh-wrapped p2wkh.
	case u.Upsbt.Inputs[inIndex].RedeemScript != nil:
		// We only need to decide if the input is witness, and we don't
		// rely on the witnessutxo/nonwitnessutxo in the PSBT, instead
		// we check the redeemScript content.
		if txscript.IsWitnessProgram(redeemScript) {
			if u.Upsbt.Inputs[inIndex].WitnessUtxo == nil {
				err := nonWitnessToWitness(u.Upsbt, inIndex)
				if err != nil {
					return SignInvalid, err
				}
			}
		}

		// If it is not a valid witness program, we here assume that
		// the provided WitnessUtxo/NonWitnessUtxo field was correct.
		err := u.addPartialSignature(inIndex, sig, pubKey)
		if err != nil {
			return SignInvalid, err
		}

	// Case 3: Neither provided only works for native p2wkh, or non-segwit
	// non-p2sh. To check if it's segwit, check the scriptPubKey of the
	// output.
	default:
		if u.Upsbt.Inputs[inIndex].WitnessUtxo == nil {
			outIndex := u.Upsbt.UnsignedTx.TxIn[inIndex].PreviousOutPoint.Index
			script := u.Upsbt.Inputs[inIndex].NonWitnessUtxo.TxOut[outIndex].PkScript

			if txscript.IsWitnessProgram(script) {
				err := nonWitnessToWitness(u.Upsbt, inIndex)
				if err != nil {
					return SignInvalid, err
				}
			}
		}

		err := u.addPartialSignature(inIndex, sig, pubKey)
		if err != nil {
			return SignInvalid, err
		}
	}

	return SignSuccesful, nil
}

// nonWitnessToWitness extracts the TxOut from the existing NonWitnessUtxo
// field in the given PSBT input and sets it as type witness by replacing the
// NonWitnessUtxo field with a WitnessUtxo field. See
// https://github.com/bitcoin/bitcoin/pull/14197.
func nonWitnessToWitness(p *Packet, inIndex int) error {
	outIndex := p.UnsignedTx.TxIn[inIndex].PreviousOutPoint.Index
	txout := p.Inputs[inIndex].NonWitnessUtxo.TxOut[outIndex]

	// Remove the non-witness first, else sanity check will not pass:
	p.Inputs[inIndex].NonWitnessUtxo = nil
	u := Updater{
		Upsbt: p,
	}

	return u.AddInWitnessUtxo(txout, inIndex)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	"github.com/pkg/errors"
)

func TestMetricBufferIsMetricWriter(t *testing.T) {
	var i interface{} = NewMetricBuffer(1)
	if _, ok := i.(MetricWriter); !ok {
		t.Fatalf("should be a MetricWriter")
	}
}

func TestMetricBufferIsMetricReader(t *testing.T) {
	var i interface{} = NewMetricBuffer(1)
	if _, ok := i.(MetricReader); !ok {
		t.Fatalf("should be a MetricReader")
	}
}

func TestMetricBufferIsMetricCloser(t *testing.T) {
	var i interface{} = NewMetricBuffer(1)
	if _, ok := i.(MetricCloser); !ok {
		t.Fatalf("should be a MetricCloser")
	}
}

func TestMetricBuffer(t *testing.T) {
	metrics := []Metric{
		{
			Name:  "test.metric.a",
			Kind:  Gauge,
			Value: 1,
		},
		{
			Name:  "test.metric.b",
			Kind:  Gauge,
			Value: 2.5,
		},
	}
	buf := NewMetricBuffer(8)
	t.Run("write slice of metrics to buffer", func(t *testing.T) {
		err := buf.WriteMetrics(metrics)
		if err != nil {
			t.Fatal(err)
		}
	})
	for i := 0; i < len(metrics); i++ {
		t.Run(fmt.Sprintf("Read %dth ('%s') metric from buffer", i, metrics[i].Name), func(t *testing.T) {
			m, err := buf.ReadMetric()
			if err != nil {
				t.Fatal(err)
			}
			if m.Value != metrics[i].Value {
				t.Fatalf("expected metric.Value to be %d got %d", metrics[i].Value, m.Value)
			}
		})
	}
	t.Run("Read should block until closed", func(t *testing.T) {
		go func() {
			time.Sleep(100 * time.Millisecond)
			buf.Close()
		}()
		_, err := buf.ReadMetric()
		if err != EOS {
			t.Fatalf("expected EOS got: %v", err)
		}
	})
}

func TestMetricPoller(t *testing.T) {
	var r MetricReadCloser = NewMetricPoller(100*time.Millisecond, func(w MetricWriter) error {
		return w.WriteMetrics([]Metric{
			{
				Name:  "test.poller",
				Time:  time.Now(),
				Kind:  Gauge,
				Value: 1,
			},
		})

	})
	t.Run("collect 5x intervals worth of data", func(t *testing.T) {
		wait := time.After(100 * 9 * time.Millisecond)
		metrics := []Metric{}
		reading := true
		for reading {
			select {
			case <-wait:
				reading = false
			default:
				m, err := r.ReadMetric()
				if err != nil {
					t.Fatal(err)
				}
				metrics = append(metrics, m)
			}
		}
		if len(metrics) > 11 || len(metrics) < 9 {
			t.Fatal("expected to collect roughly ~10 metrics over ~1 second got %v", len(metrics))
		}
	})
	t.Run("close should end polling", func(t *testing.T) {
		go func() {
			time.Sleep(10 * time.Millisecond)
			r.Close()
		}()
		_, err := r.ReadMetric()
		if err != EOS {
			t.Fatalf("expected EOS got: %v", err)
		}
	})
}

func TestMetricPollerError(t *testing.T) {
	bang := errors.New("BANG!")
	poller := NewMetricPoller(1000*time.Millisecond, func(w MetricWriter) error {
		return bang
	})
	defer poller.Close()
	_, err := poller.ReadMetric()
	if err != bang {
		t.Fatalf("expected poller to return error 'BANG!' got: %v", err)
	}
}

func TestCopyMetrics(t *testing.T) {
	inp := []Metric{
		{Name: "test.a"},
		{Name: "test.b"},
		{Name: "test.c"},
	}
	src := NewMetricBuffer(8)
	go func() {
		if err := src.WriteMetrics(inp); err != nil {
			t.Error(err)
		}
		src.Close()
	}()

	dst := NewMetricBuffer(8)
	go func() {
		if err := CopyMetrics(dst, src); err != nil {
			t.Error(err)
		}
	}()

	out := []Metric{}
	func() {
		for {
			m, err := dst.ReadMetric()
			if err == EOS {
				return
			} else if err != nil {
				t.Fatal(err)
			}
			out = append(out, m)
		}
	}()

	if !reflect.DeepEqual(inp, out) {
		t.Fatalf("expected src == dst, got: %v", out)
	}
}

func TestMultiMetricReader(t *testing.T) {
	buffers := []*MetricBuffer{
		NewMetricBuffer(4),
		NewMetricBuffer(4),
		NewMetricBuffer(4),
	}
	readers := make([]MetricReader, len(buffers))
	for i := 0; i < len(readers); i++ {
		readers[i] = buffers[i]
	}
	multi := NewMultiMetricReader(readers...)

	expectedError := fmt.Errorf("BANG!")
	buffers[0].events <- event{
		metric: Metric{Name: "should not be seen"},
		err:    expectedError,
	}

	for i := 0; i < len(buffers); i++ {
		buffers[i].WriteMetrics([]Metric{
			{Name: fmt.Sprintf("test.multi.buf%d", i)},
		})
	}

	metrics := map[string]bool{}
	errors := []error{}
	for i := 0; i < len(buffers)+1; i++ {
		m, err := multi.ReadMetric()
		if err != nil {
			errors = append(errors, err)
		} else {
			metrics[m.Name] = true
		}
	}

	t.Run("reading", func(t *testing.T) {
		if !reflect.DeepEqual(errors, []error{expectedError}) {
			t.Fatalf("expected error(s) %q got: %q", expectedError, errors)
		}

		if len(metrics) != len(buffers) {
			t.Fatalf("expected to read %d metrics out from the MultiReader got: %v", len(buffers), len(metrics))
		}

		for i := 0; i < len(buffers); i++ {
			name := fmt.Sprintf("test.multi.buf%d", i)
			t.Run(name, func(t *testing.T) {
				if !metrics[name] {
					t.Fatalf("%s metric not read from MultiReader", name)
				}
			})
		}
	})

	t.Run("closing", func(t *testing.T) {
		multi.Close()
		if _, err := multi.ReadMetric(); err != EOS {
			t.Fatal("expected MultReader to be closed")
		}
		for i := 0; i < len(buffers); i++ {
			if _, err := buffers[i].ReadMetric(); err != EOS {
				t.Fatalf("expected buffer%d to be closed", i)
			}
		}
	})
}
<commit_msg>Fix typo in tools/metrics test<commit_after>package main

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	"github.com/pkg/errors"
)

func TestMetricBufferIsMetricWriter(t *testing.T) {
	var i interface{} = NewMetricBuffer(1)
	if _, ok := i.(MetricWriter); !ok {
		t.Fatalf("should be a MetricWriter")
	}
}

func TestMetricBufferIsMetricReader(t *testing.T) {
	var i interface{} = NewMetricBuffer(1)
	if _, ok := i.(MetricReader); !ok {
		t.Fatalf("should be a MetricReader")
	}
}

func TestMetricBufferIsMetricCloser(t *testing.T) {
	var i interface{} = NewMetricBuffer(1)
	if _, ok := i.(MetricCloser); !ok {
		t.Fatalf("should be a MetricCloser")
	}
}

func TestMetricBuffer(t *testing.T) {
	metrics := []Metric{
		{
			Name:  "test.metric.a",
			Kind:  Gauge,
			Value: 1,
		},
		{
			Name:  "test.metric.b",
			Kind:  Gauge,
			Value: 2.5,
		},
	}
	buf := NewMetricBuffer(8)
	t.Run("write slice of metrics to buffer", func(t *testing.T) {
		err := buf.WriteMetrics(metrics)
		if err != nil {
			t.Fatal(err)
		}
	})
	for i := 0; i < len(metrics); i++ {
		t.Run(fmt.Sprintf("Read %dth ('%s') metric from buffer", i, metrics[i].Name), func(t *testing.T) {
			m, err := buf.ReadMetric()
			if err != nil {
				t.Fatal(err)
			}
			if m.Value != metrics[i].Value {
				t.Fatalf("expected metric.Value to be %d got %d", metrics[i].Value, m.Value)
			}
		})
	}
	t.Run("Read should block until closed", func(t *testing.T) {
		go func() {
			time.Sleep(100 * time.Millisecond)
			buf.Close()
		}()
		_, err := buf.ReadMetric()
		if err != EOS {
			t.Fatalf("expected EOS got: %v", err)
		}
	})
}

func TestMetricPoller(t *testing.T) {
	var r MetricReadCloser = NewMetricPoller(100*time.Millisecond, func(w MetricWriter) error {
		return w.WriteMetrics([]Metric{
			{
				Name:  "test.poller",
				Time:  time.Now(),
				Kind:  Gauge,
				Value: 1,
			},
		})

	})
	t.Run("collect 5x intervals worth of data", func(t *testing.T) {
		wait := time.After(100 * 9 * time.Millisecond)
		metrics := []Metric{}
		reading := true
		for reading {
			select {
			case <-wait:
				reading = false
			default:
				m, err := r.ReadMetric()
				if err != nil {
					t.Fatal(err)
				}
				metrics = append(metrics, m)
			}
		}
		if len(metrics) > 11 || len(metrics) < 9 {
			t.Fatal("expected to collect roughly ~10 metrics over ~1 second got %v", len(metrics))
		}
	})
	t.Run("close should end polling", func(t *testing.T) {
		go func() {
			time.Sleep(10 * time.Millisecond)
			r.Close()
		}()
		_, err := r.ReadMetric()
		if err != EOS {
			t.Fatalf("expected EOS got: %v", err)
		}
	})
}

func TestMetricPollerError(t *testing.T) {
	bang := errors.New("BANG!")
	poller := NewMetricPoller(1000*time.Millisecond, func(w MetricWriter) error {
		return bang
	})
	defer poller.Close()
	_, err := poller.ReadMetric()
	if err != bang {
		t.Fatalf("expected poller to return error 'BANG!' got: %v", err)
	}
}

func TestCopyMetrics(t *testing.T) {
	inp := []Metric{
		{Name: "test.a"},
		{Name: "test.b"},
		{Name: "test.c"},
	}
	src := NewMetricBuffer(8)
	go func() {
		if err := src.WriteMetrics(inp); err != nil {
			t.Error(err)
		}
		src.Close()
	}()

	dst := NewMetricBuffer(8)
	go func() {
		if err := CopyMetrics(dst, src); err != nil {
			t.Error(err)
		}
	}()

	out := []Metric{}
	func() {
		for {
			m, err := dst.ReadMetric()
			if err == EOS {
				return
			} else if err != nil {
				t.Fatal(err)
			}
			out = append(out, m)
		}
	}()

	if !reflect.DeepEqual(inp, out) {
		t.Fatalf("expected src == dst, got: %v", out)
	}
}

func TestMultiMetricReader(t *testing.T) {
	buffers := []*MetricBuffer{
		NewMetricBuffer(4),
		NewMetricBuffer(4),
		NewMetricBuffer(4),
	}
	readers := make([]MetricReader, len(buffers))
	for i := 0; i < len(readers); i++ {
		readers[i] = buffers[i]
	}
	multi := NewMultiMetricReader(readers...)

	expectedError := fmt.Errorf("BANG!")
	buffers[0].events <- event{
		metric: Metric{Name: "should not be seen"},
		err:    expectedError,
	}

	for i := 0; i < len(buffers); i++ {
		buffers[i].WriteMetrics([]Metric{
			{Name: fmt.Sprintf("test.multi.buf%d", i)},
		})
	}

	metrics := map[string]bool{}
	errors := []error{}
	for i := 0; i < len(buffers)+1; i++ {
		m, err := multi.ReadMetric()
		if err != nil {
			errors = append(errors, err)
		} else {
			metrics[m.Name] = true
		}
	}

	t.Run("reading", func(t *testing.T) {
		if !reflect.DeepEqual(errors, []error{expectedError}) {
			t.Fatalf("expected error(s) %q got: %q", expectedError, errors)
		}

		if len(metrics) != len(buffers) {
			t.Fatalf("expected to read %d metrics out from the MultiReader got: %v", len(buffers), len(metrics))
		}

		for i := 0; i < len(buffers); i++ {
			name := fmt.Sprintf("test.multi.buf%d", i)
			t.Run(name, func(t *testing.T) {
				if !metrics[name] {
					t.Fatalf("%s metric not read from MultiReader", name)
				}
			})
		}
	})

	t.Run("closing", func(t *testing.T) {
		multi.Close()
		if _, err := multi.ReadMetric(); err != EOS {
			t.Fatal("expected MultiReader to be closed")
		}
		for i := 0; i < len(buffers); i++ {
			if _, err := buffers[i].ReadMetric(); err != EOS {
				t.Fatalf("expected buffer%d to be closed", i)
			}
		}
	})
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"text/tabwriter"

	"github.com/coreos-inc/updatectl/client/update/v1"
	"github.com/coreos-inc/updatectl/third_party/github.com/codegangsta/cli"
)

func AppCommands() []cli.Command {
	return []cli.Command{
		{
			Name:  "list-apps",
			Usage: "list-apps",
			Description: `List all of the apps that exist including their label,
token and update state.`,
			Action: handle(listApps),
		},
		{
			Name:        "update-app",
			Usage:       "update-app <appId> <channelId> <label>",
			Description: `Update an app given a label.`,
			Action:      handle(updateApp),
		},
		{
			Name:        "delete-app",
			Usage:       "delete-app <appId>",
			Description: `Delete a app given an id.`,
			Action:      handle(deleteApp),
		},
	}
}

func formatApp(app *update.App) string {
	return fmt.Sprintf("%s\t%s\t%s\n", app.Id, app.Label, app.Description)
}

func listApps(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	listCall := service.App.List()
	list, err := listCall.Do()

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

	fmt.Fprintln(out, "Id\tLabel\tDescription")
	for _, app := range list.Items {
		fmt.Fprintf(out, "%s", formatApp(app))
	}

	out.Flush()
}

func updateApp(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	args := c.Args()

	if len(args) != 3 {
		cli.ShowCommandHelp(c, "update-app")
		os.Exit(1)
	}
	appReq := &update.AppUpdateReq{Label: args[1], Description: args[2]}
	call := service.App.Update(args[0], appReq)
	app, err := call.Do()

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

	fmt.Fprintf(out, "%s", formatApp(app))

	out.Flush()
}

func deleteApp(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	args := c.Args()

	if len(args) != 1 {
		fmt.Println("App token is required")
		os.Exit(1)
	}

	call := service.App.Delete(args[0])
	app, err := call.Do()

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

	fmt.Fprintf(out, "%s", formatApp(app))

	out.Flush()
}
<commit_msg>feat(create-app): add create-app alias for update-app<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"text/tabwriter"

	"github.com/coreos-inc/updatectl/third_party/code.google.com/p/go-uuid/uuid"
	"github.com/coreos-inc/updatectl/client/update/v1"
	"github.com/coreos-inc/updatectl/third_party/github.com/codegangsta/cli"
)

func AppCommands() []cli.Command {
	return []cli.Command{
		{
			Name:  "list-apps",
			Usage: "list-apps",
			Description: `List all of the apps that exist including their label,
token and update state.`,
			Action: handle(listApps),
		},
		{
			Name:        "create-app",
			Usage:       "create-app [<appId>] <label> <description>",
			Description: `Create a new application. If appId is not provided it will be created randomly.`,
			Action:      handle(createApp),
		},
		{
			Name:        "update-app",
			Usage:       "update-app <appId> <label> <description>",
			Description: `Update an app given a label.`,
			Action:      handle(updateApp),
		},
		{
			Name:        "delete-app",
			Usage:       "delete-app <appId>",
			Description: `Delete a app given an id.`,
			Action:      handle(deleteApp),
		},
	}
}

func formatApp(app *update.App) string {
	return fmt.Sprintf("%s\t%s\t%s\n", app.Id, app.Label, app.Description)
}

func listApps(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	listCall := service.App.List()
	list, err := listCall.Do()

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

	fmt.Fprintln(out, "Id\tLabel\tDescription")
	for _, app := range list.Items {
		fmt.Fprintf(out, "%s", formatApp(app))
	}

	out.Flush()
}

func createApp(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	args := c.Args()

	if !(len(args) == 2 || len(args) == 3) {
		cli.ShowCommandHelp(c, "create-app")
		os.Exit(1)
	}

	if len(args) == 2 {
		appId := uuid.New()
		args = append([]string{appId}, args...)
	}

	updateAppHelper(c, service, out, args)
}

func updateApp(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	args := c.Args()
	if len(args) != 3 {
		cli.ShowCommandHelp(c, "update-app")
		os.Exit(1)
	}
	updateAppHelper(c, service, out, args)
}

func updateAppHelper(c *cli.Context, service *update.Service, out *tabwriter.Writer, args []string) {
	appReq := &update.AppUpdateReq{Label: args[1], Description: args[2]}
	call := service.App.Update(args[0], appReq)
	app, err := call.Do()

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

	fmt.Fprintf(out, "%s", formatApp(app))

	out.Flush()
}

func deleteApp(c *cli.Context, service *update.Service, out *tabwriter.Writer) {
	args := c.Args()

	if len(args) != 1 {
		fmt.Println("App token is required")
		os.Exit(1)
	}

	call := service.App.Delete(args[0])
	app, err := call.Do()

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

	fmt.Fprintf(out, "%s", formatApp(app))

	out.Flush()
}
<|endoftext|>


<commit_before>package tlock

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/siddontang/goredis"
)

var errLockTimeout = errors.New("lock timeout")

type App struct {
	m sync.Mutex

	wg sync.WaitGroup

	httpListener net.Listener
	respListener net.Listener

	keyLockerGroup  *KeyLockerGroup
	pathLockerGroup *PathLockerGroup

	locksMutex sync.Mutex
	locks      map[uint64]*lockInfo

	lockIDCounter uint32
}

type lockInfo struct {
	id         uint64
	names      []string
	tp         string
	createTime time.Time
}

func newLockInfo(id uint64, tp string, names []string) *lockInfo {
	l := new(lockInfo)

	l.id = id
	l.names = names
	l.tp = tp
	l.createTime = time.Now()

	return l
}

type lockInfos []*lockInfo

func (s lockInfos) Len() int {
	return len(s)
}

func (s lockInfos) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s lockInfos) Less(i, j int) bool {
	return s[i].id < s[j].id
}

func NewApp() *App {
	a := new(App)

	a.keyLockerGroup = NewKeyLockerGroup()
	a.pathLockerGroup = NewPathLockerGroup()

	a.locks = make(map[uint64]*lockInfo, 1024)

	return a
}

func (a *App) StartHTTP(addr string) error {
	a.m.Lock()
	defer a.m.Unlock()

	var err error
	a.httpListener, err = net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()

		mux := http.NewServeMux()
		mux.Handle("/lock", a.newLockHandler())

		http.Serve(a.httpListener, mux)

	}()
	return nil
}

func (a *App) StartRESP(addr string) error {
	a.m.Lock()
	defer a.m.Unlock()

	var err error
	a.respListener, err = net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()

		for {
			conn, err := a.respListener.Accept()
			if err != nil {
				return
			}

			go a.handleRESP(conn)
		}

	}()
	return nil
}

func (a *App) Close() {
	a.m.Lock()
	defer a.m.Unlock()

	if a.httpListener != nil {
		a.httpListener.Close()
	}

	if a.respListener != nil {
		a.respListener.Close()
	}

	a.wg.Wait()
}

func (a *App) HTTPAddr() net.Addr {
	if a.httpListener == nil {
		return nil
	} else {
		return a.httpListener.Addr()
	}
}

func (a *App) RESPAddr() net.Addr {
	if a.respListener == nil {
		return nil
	} else {
		return a.respListener.Addr()
	}
}

func (a *App) genLockID() uint64 {
	//todo, optimize later
	id := uint64(time.Now().Unix())
	c := uint64(atomic.AddUint32(&a.lockIDCounter, 1))
	return id<<32 | c
}

// Lock and returns a lock id, you must use this id to unlock
func (a *App) Lock(tp string, names []string) (uint64, error) {
	id, err := a.LockTimeout(tp, InfiniteTimeout, names)
	return id, err
}

// Lock with timeout and returns a lock id, you must use this id to unlock
func (a *App) LockTimeout(tp string, timeout time.Duration, names []string) (uint64, error) {
	if len(names) == 0 {
		return 0, fmt.Errorf("empty lock names")
	}

	var b bool
	var err error
	tp = strings.ToLower(tp)
	switch tp {
	case "key":
		b, err = a.keyLockerGroup.LockTimeout(timeout, names...), nil
	case "path":
		b, err = a.pathLockerGroup.LockTimeout(timeout, names...), nil
	default:
		return 0, fmt.Errorf("invalid lock type %s", tp)
	}
	if !b {
		return 0, errLockTimeout
	} else if err != nil {
		return 0, err
	}

	id := a.genLockID()
	l := newLockInfo(id, tp, names)

	a.locksMutex.Lock()
	a.locks[id] = l
	a.locksMutex.Unlock()
	return id, nil
}

func (a *App) Unlock(id uint64) error {
	if id == 0 {
		return fmt.Errorf("empty lock names")
	}

	a.locksMutex.Lock()
	l, ok := a.locks[id]
	delete(a.locks, id)
	a.locksMutex.Unlock()

	if !ok {
		return nil
	}

	switch l.tp {
	case "key":
		a.keyLockerGroup.Unlock(l.names...)
	case "path":
		a.pathLockerGroup.Unlock(l.names...)
	default:
		return fmt.Errorf("invalid lock type %s", l.tp)
	}

	return nil
}

const timeFormat string = "2006-01-02 15:04:05"

func (a *App) dumpLockNames() []byte {
	var buf bytes.Buffer

	keyLocks := make(lockInfos, 0, 1024)
	pathLocks := make(lockInfos, 0, 1024)

	a.locksMutex.Lock()
	for _, l := range a.locks {
		if l.tp == "key" {
			keyLocks = append(keyLocks, l)
		} else {
			pathLocks = append(pathLocks, l)
		}
	}
	a.locksMutex.Unlock()

	sort.Sort(keyLocks)
	sort.Sort(pathLocks)

	buf.WriteString("key lock:\n")
	for _, l := range keyLocks {
		buf.WriteString(fmt.Sprintf("%d %v\t%s\n", l.id, l.names, l.createTime.Format(timeFormat)))
	}

	buf.WriteString("\npath lock:\n")
	for _, l := range pathLocks {
		buf.WriteString(fmt.Sprintf("%d %v\t%s\n", l.id, l.names, l.createTime.Format(timeFormat)))
	}

	return buf.Bytes()
}

// lock name1, name2, ... [TYPE key] [TIMEOUT 60]
// unlock id
func (a *App) handleRESP(c net.Conn) {
	conn, err := goredis.NewConn(c)
	if err != nil {
		c.Close()
		return
	}

	grapLockIDs := make(map[uint64]struct{})

	defer func() {
		conn.Close()
		for id, _ := range grapLockIDs {
			a.Unlock(id)
		}
	}()

	for {
		args, err := conn.ReceiveRequest()
		if err != nil {
			return
		}
		if len(args) < 1 {
			conn.SendValue(fmt.Errorf("empty command"))
			continue
		}

		cmd := strings.ToUpper(string(args[0]))
		args = args[1:]
		switch cmd {
		case "LOCK":
			tp, names, timeout, err := a.parseRESPLock(args)
			if err != nil {
				conn.SendValue(err)
			} else {
				id, err := a.LockTimeout(tp, timeout, names)
				if err != nil {
					conn.SendValue(err)
				} else {
					grapLockIDs[id] = struct{}{}
					conn.SendValue([]byte(strconv.FormatUint(id, 10)))
				}
			}
		case "UNLOCK":
			id, err := a.parseRESPUnlock(args)
			if err != nil {
				conn.SendValue(err)
			} else {
				err = a.Unlock(id)
				if err != nil {
					conn.SendValue(err)
				} else {
					delete(grapLockIDs, id)
					conn.SendValue("OK")
				}
			}
		default:
			conn.SendValue(fmt.Errorf("invalid command %s", cmd))
		}
	}
}

func (a *App) parseRESPLock(args [][]byte) (tp string, names []string, timeout time.Duration, err error) {
	tp = "key"
	timeout = 60 * time.Second

	names = make([]string, 0, len(args))

	for i := 0; i < len(args); i++ {
		arg := string(args[i])
		s := strings.ToUpper(arg)
		if s == "TYPE" && i < len(args) {
			tp = string(args[i+1])
			i++
		} else if s == "TIMEOUT" && i < len(args) {
			var t uint64
			t, err = strconv.ParseUint(string(args[i+1]), 10, 64)
			if err != nil {
				return
			}
			if t > 60 || t == 0 {
				t = 60
			}

			timeout = time.Duration(t) * time.Second
			i++
		} else {
			names = append(names, arg)
		}
	}
	return
}

func (a *App) parseRESPUnlock(args [][]byte) (id uint64, err error) {
	if len(args) != 1 {
		return 0, fmt.Errorf("empty unlock id")
	}

	return strconv.ParseUint(string(args[0]), 10, 64)
}

type lockHandler struct {
	a *App
}

func (a *App) newLockHandler() *lockHandler {
	h := new(lockHandler)
	h.a = a

	return h
}

// Lock:   Post/Put /lock?names=a,b,c&timeout=10&type=key return a lock id
// Unlock: Delete   /lock?id=lockid
// For HTTP, the default and maximum timeout is 60s
// Lock type supports key and path, the default is key
// List locks: Get  /lock
func (h *lockHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		buf := h.a.dumpLockNames()
		w.Header().Set("Content-Type", "text/plain")
		w.Write(buf)
		return
	case "POST", "PUT":
		names := strings.Split(r.FormValue("names"), ",")
		if len(names) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("empty lock names"))
			return
		}

		timeout, _ := strconv.Atoi(r.FormValue("timeout"))

		if timeout <= 0 || timeout > 60 {
			timeout = 60
		}
		tp := strings.ToLower(r.FormValue("type"))
		if len(tp) == 0 {
			tp = "key"
		}

		id, err := h.a.LockTimeout(tp, time.Duration(timeout)*time.Second, names)
		if err != nil && err != errLockTimeout {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(err.Error()))
		} else if err == errLockTimeout {
			w.WriteHeader(http.StatusRequestTimeout)
			w.Write([]byte("Lock timeout"))
		} else {
			w.WriteHeader(http.StatusOK)
			w.Write([]byte(strconv.FormatUint(id, 10)))
		}
	case "DELETE":
		id, err := strconv.ParseUint(r.FormValue("id"), 10, 64)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(err.Error()))
			return
		}

		err = h.a.Unlock(id)

		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(err.Error()))
		} else {
			w.WriteHeader(http.StatusOK)
		}
	default:
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
}
<commit_msg>update timeout handling<commit_after>package tlock

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/siddontang/goredis"
)

var errLockTimeout = errors.New("lock timeout")

type App struct {
	m sync.Mutex

	wg sync.WaitGroup

	httpListener net.Listener
	respListener net.Listener

	keyLockerGroup  *KeyLockerGroup
	pathLockerGroup *PathLockerGroup

	locksMutex sync.Mutex
	locks      map[uint64]*lockInfo

	lockIDCounter uint32
}

type lockInfo struct {
	id         uint64
	names      []string
	tp         string
	createTime time.Time
}

func newLockInfo(id uint64, tp string, names []string) *lockInfo {
	l := new(lockInfo)

	l.id = id
	l.names = names
	l.tp = tp
	l.createTime = time.Now()

	return l
}

type lockInfos []*lockInfo

func (s lockInfos) Len() int {
	return len(s)
}

func (s lockInfos) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s lockInfos) Less(i, j int) bool {
	return s[i].id < s[j].id
}

func NewApp() *App {
	a := new(App)

	a.keyLockerGroup = NewKeyLockerGroup()
	a.pathLockerGroup = NewPathLockerGroup()

	a.locks = make(map[uint64]*lockInfo, 1024)

	return a
}

func (a *App) StartHTTP(addr string) error {
	a.m.Lock()
	defer a.m.Unlock()

	var err error
	a.httpListener, err = net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()

		mux := http.NewServeMux()
		mux.Handle("/lock", a.newLockHandler())

		http.Serve(a.httpListener, mux)

	}()
	return nil
}

func (a *App) StartRESP(addr string) error {
	a.m.Lock()
	defer a.m.Unlock()

	var err error
	a.respListener, err = net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()

		for {
			conn, err := a.respListener.Accept()
			if err != nil {
				return
			}

			go a.handleRESP(conn)
		}

	}()
	return nil
}

func (a *App) Close() {
	a.m.Lock()
	defer a.m.Unlock()

	if a.httpListener != nil {
		a.httpListener.Close()
	}

	if a.respListener != nil {
		a.respListener.Close()
	}

	a.wg.Wait()
}

func (a *App) HTTPAddr() net.Addr {
	if a.httpListener == nil {
		return nil
	} else {
		return a.httpListener.Addr()
	}
}

func (a *App) RESPAddr() net.Addr {
	if a.respListener == nil {
		return nil
	} else {
		return a.respListener.Addr()
	}
}

func (a *App) genLockID() uint64 {
	//todo, optimize later
	id := uint64(time.Now().Unix())
	c := uint64(atomic.AddUint32(&a.lockIDCounter, 1))
	return id<<32 | c
}

// Lock and returns a lock id, you must use this id to unlock
func (a *App) Lock(tp string, names []string) (uint64, error) {
	id, err := a.LockTimeout(tp, InfiniteTimeout, names)
	return id, err
}

// Lock with timeout and returns a lock id, you must use this id to unlock
func (a *App) LockTimeout(tp string, timeout time.Duration, names []string) (uint64, error) {
	if len(names) == 0 {
		return 0, fmt.Errorf("empty lock names")
	}

	var b bool
	var err error
	tp = strings.ToLower(tp)
	switch tp {
	case "key":
		b, err = a.keyLockerGroup.LockTimeout(timeout, names...), nil
	case "path":
		b, err = a.pathLockerGroup.LockTimeout(timeout, names...), nil
	default:
		return 0, fmt.Errorf("invalid lock type %s", tp)
	}
	if !b {
		return 0, errLockTimeout
	} else if err != nil {
		return 0, err
	}

	id := a.genLockID()
	l := newLockInfo(id, tp, names)

	a.locksMutex.Lock()
	a.locks[id] = l
	a.locksMutex.Unlock()
	return id, nil
}

func (a *App) Unlock(id uint64) error {
	if id == 0 {
		return fmt.Errorf("empty lock names")
	}

	a.locksMutex.Lock()
	l, ok := a.locks[id]
	delete(a.locks, id)
	a.locksMutex.Unlock()

	if !ok {
		return nil
	}

	switch l.tp {
	case "key":
		a.keyLockerGroup.Unlock(l.names...)
	case "path":
		a.pathLockerGroup.Unlock(l.names...)
	default:
		return fmt.Errorf("invalid lock type %s", l.tp)
	}

	return nil
}

const timeFormat string = "2006-01-02 15:04:05"

func (a *App) dumpLockNames() []byte {
	var buf bytes.Buffer

	keyLocks := make(lockInfos, 0, 1024)
	pathLocks := make(lockInfos, 0, 1024)

	a.locksMutex.Lock()
	for _, l := range a.locks {
		if l.tp == "key" {
			keyLocks = append(keyLocks, l)
		} else {
			pathLocks = append(pathLocks, l)
		}
	}
	a.locksMutex.Unlock()

	sort.Sort(keyLocks)
	sort.Sort(pathLocks)

	buf.WriteString("key lock:\n")
	for _, l := range keyLocks {
		buf.WriteString(fmt.Sprintf("%d %v\t%s\n", l.id, l.names, l.createTime.Format(timeFormat)))
	}

	buf.WriteString("\npath lock:\n")
	for _, l := range pathLocks {
		buf.WriteString(fmt.Sprintf("%d %v\t%s\n", l.id, l.names, l.createTime.Format(timeFormat)))
	}

	return buf.Bytes()
}

// lock name1, name2, ... [TYPE key] [TIMEOUT 60]
// unlock id
func (a *App) handleRESP(c net.Conn) {
	conn, err := goredis.NewConn(c)
	if err != nil {
		c.Close()
		return
	}

	grapLockIDs := make(map[uint64]struct{})

	defer func() {
		conn.Close()
		for id, _ := range grapLockIDs {
			a.Unlock(id)
		}
	}()

	for {
		args, err := conn.ReceiveRequest()
		if err != nil {
			return
		}
		if len(args) < 1 {
			conn.SendValue(fmt.Errorf("empty command"))
			continue
		}

		cmd := strings.ToUpper(string(args[0]))
		args = args[1:]
		switch cmd {
		case "LOCK":
			tp, names, timeout, err := a.parseRESPLock(args)
			if err != nil {
				conn.SendValue(err)
			} else {
				id, err := a.LockTimeout(tp, timeout, names)
				if err != nil {
					conn.SendValue(err)
				} else {
					grapLockIDs[id] = struct{}{}
					conn.SendValue([]byte(strconv.FormatUint(id, 10)))
				}
			}
		case "UNLOCK":
			id, err := a.parseRESPUnlock(args)
			if err != nil {
				conn.SendValue(err)
			} else {
				err = a.Unlock(id)
				if err != nil {
					conn.SendValue(err)
				} else {
					delete(grapLockIDs, id)
					conn.SendValue("OK")
				}
			}
		default:
			conn.SendValue(fmt.Errorf("invalid command %s", cmd))
		}
	}
}

func (a *App) parseRESPLock(args [][]byte) (tp string, names []string, timeout time.Duration, err error) {
	tp = "key"
	timeout = 60 * time.Second

	names = make([]string, 0, len(args))

	for i := 0; i < len(args); i++ {
		arg := string(args[i])
		s := strings.ToUpper(arg)
		if s == "TYPE" && i < len(args) {
			tp = strings.ToLower(string(args[i+1]))
			i++
		} else if s == "TIMEOUT" && i < len(args) {
			var t uint64
			t, err = strconv.ParseUint(string(args[i+1]), 10, 64)
			if err != nil {
				return
			}
			if t == 0 {
				t = 60
			}

			timeout = time.Duration(t) * time.Second
			i++
		} else {
			names = append(names, arg)
		}
	}
	return
}

func (a *App) parseRESPUnlock(args [][]byte) (id uint64, err error) {
	if len(args) != 1 {
		return 0, fmt.Errorf("empty unlock id")
	}

	return strconv.ParseUint(string(args[0]), 10, 64)
}

type lockHandler struct {
	a *App
}

func (a *App) newLockHandler() *lockHandler {
	h := new(lockHandler)
	h.a = a

	return h
}

// Lock:   Post/Put /lock?names=a,b,c&timeout=10&type=key return a lock id
// Unlock: Delete   /lock?id=lockid
// For HTTP, the default and maximum timeout is 60s
// Lock type supports key and path, the default is key
// List locks: Get  /lock
func (h *lockHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		buf := h.a.dumpLockNames()
		w.Header().Set("Content-Type", "text/plain")
		w.Write(buf)
		return
	case "POST", "PUT":
		names := strings.Split(r.FormValue("names"), ",")
		if len(names) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("empty lock names"))
			return
		}

		timeout, _ := strconv.Atoi(r.FormValue("timeout"))

		if timeout <= 0 {
			timeout = 60
		}
		tp := strings.ToLower(r.FormValue("type"))
		if len(tp) == 0 {
			tp = "key"
		}

		id, err := h.a.LockTimeout(tp, time.Duration(timeout)*time.Second, names)
		if err != nil && err != errLockTimeout {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(err.Error()))
		} else if err == errLockTimeout {
			w.WriteHeader(http.StatusRequestTimeout)
			w.Write([]byte("Lock timeout"))
		} else {
			w.WriteHeader(http.StatusOK)
			w.Write([]byte(strconv.FormatUint(id, 10)))
		}
	case "DELETE":
		id, err := strconv.ParseUint(r.FormValue("id"), 10, 64)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(err.Error()))
			return
		}

		err = h.a.Unlock(id)

		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(err.Error()))
		} else {
			w.WriteHeader(http.StatusOK)
		}
	default:
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
}
<|endoftext|>