<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"

        "github.com/johnny-morrice/ctrl"
        "github.com/gorilla/mux"
        "github.com/gorilla/sessions"
)

// testcase is a controller that runs on the given port.
type testcase struct {
	port     int
	set      *testset
	commands chan<- command
	hostname string
}

func (tcase *testcase) handler() http.Handler {
	r := mux.NewRouter()
	r.Handle("/api/test/{resultset}/ping", tcase)

        // Use CORS to allow all origins.
        handler := newcorshandler(r, fmt.Sprintf("http://%v", tcase.hostname))

        return handler
}

func (tc *testcase) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Forward declaration for goto.
	ping := ping{}
	cmd := command{}
        reply := pingreply{}
        var session *sessions.Session
        var userid int

	c := ctrl.New(w, r)

	rset, err := resultsetparam(c)

	if err != nil {
		goto ERROR
	}

        session, err = getsession(r)

        if err != nil {
                log.Printf("Session error: %v", err)
        }

        userid = getsessionuser(session)

        // Send command to main loop.
	ping.port = tc.port
	ping.set = rset
        ping.reply = make(chan pingreply)
        ping.userid = userid

	cmd.ctype = _PING
	cmd.ping = ping

	tc.commands<- cmd

	// Receive command from main loop.
	reply = <-ping.reply
        err = reply.err

        if err != nil {
                goto ERROR
        }

	err = c.ServeJson(true)

	if err == nil {
		return
	}

ERROR:
	if err != nil {
		log.Printf("Error in testcase handler: %v", err)

		c.InternalError()
	}
}

// phapi is the API controller.
type phapi struct {
	commands chan<- command
}

func (api *phapi) getresults(w http.ResponseWriter, r *http.Request) {
	// Forward declare variables because we are using goto
	q := query{}
	cmd := command{}
        reply := queryreply{}
        var session *sessions.Session
        var userid int

	if debug {
		fmt.Fprintln(os.Stderr, "getresults")
	}

	c := ctrl.New(w, r)

	resultset, err := resultsetparam(c)

	if err != nil {
		goto ERROR
	}

        session, err = getsession(r)

        if err != nil {
		log.Printf("Session error: %v", err)
        }

        userid = getsessionuser(session)

        // Send command to main loop.
	q.rset = resultset
	q.reply = make(chan queryreply)
        q.userid = userid

	cmd.ctype = _GETRESULT
	cmd.query = q

	api.commands<- cmd

        // Receive command from main loop.
	reply = <-q.reply
        err = reply.err

        if err != nil {
                goto ERROR
        }

	err = c.ServeJson(reply.badports)

	if err == nil {
		return
	}

ERROR:
	if err != nil {
		log.Printf("Error in getresults: %v", err)

		c.InternalError()
	}
}

func (api *phapi) newtest(w http.ResponseWriter, r *http.Request) {
	c := ctrl.New(w, r)

        var userid int
        var session *sessions.Session
        reg := registration{}
        cmd := command{}
        reply := regisreply{}

        packet := &LaunchData{}
	dec := json.NewDecoder(r.Body)
	err := dec.Decode(packet)

	if err != nil {
		goto ERROR
	}

        session, err = getsession(r)

        if err != nil {
		log.Printf("Session error: %v", err)
        }

        userid = getsessionuser(session)

        // Send command to main loop.
	reg.reply = make(chan regisreply)
	reg.LaunchData = *packet
        reg.userid = userid

	cmd.ctype = _NEWTEST
	cmd.reg = reg

	api.commands <- cmd

        // Receive reply from main loop.
	reply = <-reg.reply
        err = reply.err

        if err != nil {
                goto ERROR
        }

        if userid != reply.userid {
                setsessionuser(session, reply.userid)

                err = session.Save(r, w)

                if err != nil {
                        goto ERROR
                }
        }

	err = c.ServeJson(reply.scanid)

ERROR:
	if err != nil {
		log.Printf("Error serving newtest: %v", err)

		c.InternalError()
	}
}

type corshandler struct {
	handler http.Handler
	allowed map[string]struct{}
	headers []string
}

func newcorshandler(h http.Handler, origins... string) http.Handler {
	ch := corshandler{}
	ch.handler = h
	ch.headers = []string{"Content-Type"}

	ch.allowed = map[string]struct{}{}
	for _, o := range origins {
		ch.allowed[o] = struct{}{}
	}

	return ch
}

func (ch corshandler) cors(w http.ResponseWriter, req *http.Request) {
	originheads := req.Header["Origin"]

	var origin string
	if len(originheads) == 1 {
		origin = originheads[0]
	} else {
		if trace {
                	log.Printf("Bad Origin header: %v", originheads)

		}

		return
        }

        _, any := ch.allowed["*"];
        _, ok := ch.allowed[origin];

        if any || ok {
		w.Header()["Access-Control-Allow-Credentials"] = []string{"true"}
		w.Header()["Access-Control-Allow-Origin"] = []string{origin}
		w.Header()["Access-Control-Allow-Headers"] = ch.headers
	}
}

func (ch corshandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	ch.cors(w, req)

	if req.Method == "OPTIONS" {
		w.WriteHeader(http.StatusOK)
	} else {
		ch.handler.ServeHTTP(w, req)
        }
}

func resultsetparam(c ctrl.C) (uint64, error) {
	resultset, err := c.GetMuxVar("resultset")

	if err != nil {
		return 0, err
	}

	setid, err := strconv.ParseUint(resultset, 10, 64)

	if err != nil {
		return 0, err
	}

	return setid, nil
}

var globalSessionStore *sessions.CookieStore

func getsession(r *http.Request) (*sessions.Session, error) {
        return globalSessionStore.Get(r, "session-name")
}

func getsessionuser(session *sessions.Session) int {
        maybeid, ok := session.Values["user"]

	if debug {
		log.Printf("session.Values: %v", session.Values)
		log.Printf("user id is: %v", maybeid)
	}

        if ok {
                var id int
                id, ok = maybeid.(int)

                if ok {
                    return id
                }
        }

        return nouser
}

func setsessionuser(session *sessions.Session, userid int) {
        session.Values["user"] = userid
}
<commit_msg>Remove useless newline<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"

        "github.com/johnny-morrice/ctrl"
        "github.com/gorilla/mux"
        "github.com/gorilla/sessions"
)

// testcase is a controller that runs on the given port.
type testcase struct {
	port     int
	set      *testset
	commands chan<- command
	hostname string
}

func (tcase *testcase) handler() http.Handler {
	r := mux.NewRouter()
	r.Handle("/api/test/{resultset}/ping", tcase)

        // Use CORS to allow all origins.
        handler := newcorshandler(r, fmt.Sprintf("http://%v", tcase.hostname))

        return handler
}

func (tc *testcase) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Forward declaration for goto.
	ping := ping{}
	cmd := command{}
        reply := pingreply{}
        var session *sessions.Session
        var userid int

	c := ctrl.New(w, r)

	rset, err := resultsetparam(c)

	if err != nil {
		goto ERROR
	}

        session, err = getsession(r)

        if err != nil {
                log.Printf("Session error: %v", err)
        }

        userid = getsessionuser(session)

        // Send command to main loop.
	ping.port = tc.port
	ping.set = rset
        ping.reply = make(chan pingreply)
        ping.userid = userid

	cmd.ctype = _PING
	cmd.ping = ping

	tc.commands<- cmd

	// Receive command from main loop.
	reply = <-ping.reply
        err = reply.err

        if err != nil {
                goto ERROR
        }

	err = c.ServeJson(true)

	if err == nil {
		return
	}

ERROR:
	if err != nil {
		log.Printf("Error in testcase handler: %v", err)

		c.InternalError()
	}
}

// phapi is the API controller.
type phapi struct {
	commands chan<- command
}

func (api *phapi) getresults(w http.ResponseWriter, r *http.Request) {
	// Forward declare variables because we are using goto
	q := query{}
	cmd := command{}
        reply := queryreply{}
        var session *sessions.Session
        var userid int

	if debug {
		fmt.Fprintln(os.Stderr, "getresults")
	}

	c := ctrl.New(w, r)

	resultset, err := resultsetparam(c)

	if err != nil {
		goto ERROR
	}

        session, err = getsession(r)

        if err != nil {
		log.Printf("Session error: %v", err)
        }

        userid = getsessionuser(session)

        // Send command to main loop.
	q.rset = resultset
	q.reply = make(chan queryreply)
        q.userid = userid

	cmd.ctype = _GETRESULT
	cmd.query = q

	api.commands<- cmd

        // Receive command from main loop.
	reply = <-q.reply
        err = reply.err

        if err != nil {
                goto ERROR
        }

	err = c.ServeJson(reply.badports)

	if err == nil {
		return
	}

ERROR:
	if err != nil {
		log.Printf("Error in getresults: %v", err)

		c.InternalError()
	}
}

func (api *phapi) newtest(w http.ResponseWriter, r *http.Request) {
	c := ctrl.New(w, r)

        var userid int
        var session *sessions.Session
        reg := registration{}
        cmd := command{}
        reply := regisreply{}

        packet := &LaunchData{}
	dec := json.NewDecoder(r.Body)
	err := dec.Decode(packet)

	if err != nil {
		goto ERROR
	}

        session, err = getsession(r)

        if err != nil {
		log.Printf("Session error: %v", err)
        }

        userid = getsessionuser(session)

        // Send command to main loop.
	reg.reply = make(chan regisreply)
	reg.LaunchData = *packet
        reg.userid = userid

	cmd.ctype = _NEWTEST
	cmd.reg = reg

	api.commands <- cmd

        // Receive reply from main loop.
	reply = <-reg.reply
        err = reply.err

        if err != nil {
                goto ERROR
        }

        if userid != reply.userid {
                setsessionuser(session, reply.userid)

                err = session.Save(r, w)

                if err != nil {
                        goto ERROR
                }
        }

	err = c.ServeJson(reply.scanid)

ERROR:
	if err != nil {
		log.Printf("Error serving newtest: %v", err)

		c.InternalError()
	}
}

type corshandler struct {
	handler http.Handler
	allowed map[string]struct{}
	headers []string
}

func newcorshandler(h http.Handler, origins... string) http.Handler {
	ch := corshandler{}
	ch.handler = h
	ch.headers = []string{"Content-Type"}

	ch.allowed = map[string]struct{}{}
	for _, o := range origins {
		ch.allowed[o] = struct{}{}
	}

	return ch
}

func (ch corshandler) cors(w http.ResponseWriter, req *http.Request) {
	originheads := req.Header["Origin"]

	var origin string
	if len(originheads) == 1 {
		origin = originheads[0]
	} else {
		if trace {
                	log.Printf("Bad Origin header: %v", originheads)
		}

		return
        }

        _, any := ch.allowed["*"];
        _, ok := ch.allowed[origin];

        if any || ok {
		w.Header()["Access-Control-Allow-Credentials"] = []string{"true"}
		w.Header()["Access-Control-Allow-Origin"] = []string{origin}
		w.Header()["Access-Control-Allow-Headers"] = ch.headers
	}
}

func (ch corshandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	ch.cors(w, req)

	if req.Method == "OPTIONS" {
		w.WriteHeader(http.StatusOK)
	} else {
		ch.handler.ServeHTTP(w, req)
        }
}

func resultsetparam(c ctrl.C) (uint64, error) {
	resultset, err := c.GetMuxVar("resultset")

	if err != nil {
		return 0, err
	}

	setid, err := strconv.ParseUint(resultset, 10, 64)

	if err != nil {
		return 0, err
	}

	return setid, nil
}

var globalSessionStore *sessions.CookieStore

func getsession(r *http.Request) (*sessions.Session, error) {
        return globalSessionStore.Get(r, "session-name")
}

func getsessionuser(session *sessions.Session) int {
        maybeid, ok := session.Values["user"]

	if debug {
		log.Printf("session.Values: %v", session.Values)
		log.Printf("user id is: %v", maybeid)
	}

        if ok {
                var id int
                id, ok = maybeid.(int)

                if ok {
                    return id
                }
        }

        return nouser
}

func setsessionuser(session *sessions.Session, userid int) {
        session.Values["user"] = userid
}
<|endoftext|>


<commit_before>package server

import (
	"net/http"
	"path"
	"strings"

	"github.com/mholt/caddy/middleware/browse"
)

// This FileServer is adapted from the one in net/http
// by the Go authors. Some modifications have been made.
//
// License:
//
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
func FileServer(root http.FileSystem) http.Handler {
	return &fileHandler{root}
}

type fileHandler struct {
	root http.FileSystem
}

func (f *fileHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	upath := r.URL.Path
	if !strings.HasPrefix(upath, "/") {
		upath = "/" + upath
		r.URL.Path = upath
	}
	serveFile(w, r, f.root, path.Clean(upath), true)
}

// name is '/'-separated, not filepath.Separator.
func serveFile(w http.ResponseWriter, r *http.Request, fs http.FileSystem, name string, redirect bool) {
	f, err := fs.Open(name)
	if err != nil {
		// TODO expose actual error?
		http.NotFound(w, r)
		return
	}
	defer f.Close()

	d, err1 := f.Stat()
	if err1 != nil {
		// TODO expose actual error?
		http.NotFound(w, r)
		return
	}

	// use contents of an index file, if present, for directory
	if d.IsDir() {
		for _, indexPage := range browse.IndexPages {
			index := strings.TrimSuffix(name, "/") + "/" + indexPage
			ff, err := fs.Open(index)
			if err == nil {
				defer ff.Close()
				dd, err := ff.Stat()
				if err == nil {
					name = index
					d = dd
					f = ff
					break
				}
			}
		}
	}

	// Still a directory? (we didn't find an index file)
	if d.IsDir() {
		http.NotFound(w, r) // 404 instead of 403 to hide the fact that the folder exists
		return
	}

	http.ServeContent(w, r, d.Name(), d.ModTime(), f)
}
<commit_msg>Refactored file server to return errors<commit_after>package server

import (
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/mholt/caddy/middleware"
	"github.com/mholt/caddy/middleware/browse"
)

// This FileServer is adapted from the one in net/http
// by the Go authors. Some modifications have been made.
//
// License:
//
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
func FileServer(root http.FileSystem) middleware.Handler {
	return &fileHandler{root}
}

type fileHandler struct {
	root http.FileSystem
}

func (f *fileHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
	upath := r.URL.Path
	if !strings.HasPrefix(upath, "/") {
		upath = "/" + upath
		r.URL.Path = upath
	}
	return serveFile(w, r, f.root, path.Clean(upath))
}

// name is '/'-separated, not filepath.Separator.
func serveFile(w http.ResponseWriter, r *http.Request, fs http.FileSystem, name string) (int, error) {
	f, err := fs.Open(name)
	if err != nil {
		if os.IsPermission(err) {
			return http.StatusForbidden, err
		}
		return http.StatusNotFound, nil
	}
	defer f.Close()

	d, err1 := f.Stat()
	if err1 != nil {
		if os.IsPermission(err) {
			return http.StatusForbidden, err
		}
		return http.StatusNotFound, nil
	}

	// use contents of an index file, if present, for directory
	if d.IsDir() {
		for _, indexPage := range browse.IndexPages {
			index := strings.TrimSuffix(name, "/") + "/" + indexPage
			ff, err := fs.Open(index)
			if err == nil {
				defer ff.Close()
				dd, err := ff.Stat()
				if err == nil {
					name = index
					d = dd
					f = ff
					break
				}
			}
		}
	}

	// Still a directory? (we didn't find an index file)
	// Return 404 to hide the fact that the folder exists
	if d.IsDir() {
		return http.StatusNotFound, nil
	}

	// Note: Errors generated by ServeContent are written immediately
	// to the response. This usually only happens if seeking fails (rare).
	http.ServeContent(w, r, d.Name(), d.ModTime(), f)

	return http.StatusOK, nil
}
<|endoftext|>


<commit_before>package agent

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"testing"
	"time"

	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/testutil"
)

func makeHTTPServer(t *testing.T) (string, *HTTPServer) {
	return makeHTTPServerWithConfig(t, nil)
}

func makeHTTPServerWithConfig(t *testing.T, cb func(c *Config)) (string, *HTTPServer) {
	conf := nextConfig()
	if cb != nil {
		cb(conf)
	}

	dir, agent := makeAgent(t, conf)
	uiDir := filepath.Join(dir, "ui")
	if err := os.Mkdir(uiDir, 755); err != nil {
		t.Fatalf("err: %v", err)
	}
	conf.UiDir = uiDir
	servers, err := NewHTTPServers(agent, conf, agent.logOutput)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if len(servers) == 0 {
		t.Fatalf(fmt.Sprintf("Failed to make HTTP server"))
	}
	return dir, servers[0]
}

func encodeReq(obj interface{}) io.ReadCloser {
	buf := bytes.NewBuffer(nil)
	enc := json.NewEncoder(buf)
	enc.Encode(obj)
	return ioutil.NopCloser(buf)
}

func TestHTTPServer_UnixSocket(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.SkipNow()
	}

	tempDir, err := ioutil.TempDir("", "consul")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(tempDir)
	socket := filepath.Join(tempDir, "test.sock")

	dir, srv := makeHTTPServerWithConfig(t, func(c *Config) {
		c.Addresses.HTTP = "unix://" + socket
	})
	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	// Ensure the socket was created
	if _, err := os.Stat(socket); err != nil {
		t.Fatalf("err: %s", err)
	}

	// Ensure we can get a response from the socket.
	path, _ := unixSocketAddr(srv.agent.config.Addresses.HTTP)
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(_, _ string) (net.Conn, error) {
				return net.Dial("unix", path)
			},
		},
	}

	// This URL doesn't look like it makes sense, but the scheme (http://) and
	// the host (127.0.0.1) are required by the HTTP client library. In reality
	// this will just use the custom dialer and talk to the socket.
	resp, err := client.Get("http://127.0.0.1/v1/agent/self")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer resp.Body.Close()

	if body, err := ioutil.ReadAll(resp.Body); err != nil || len(body) == 0 {
		t.Fatalf("bad: %s %v", body, err)
	}
}

func TestSetIndex(t *testing.T) {
	resp := httptest.NewRecorder()
	setIndex(resp, 1000)
	header := resp.Header().Get("X-Consul-Index")
	if header != "1000" {
		t.Fatalf("Bad: %v", header)
	}
	setIndex(resp, 2000)
	if v := resp.Header()["X-Consul-Index"]; len(v) != 1 {
		t.Fatalf("bad: %#v", v)
	}
}

func TestSetKnownLeader(t *testing.T) {
	resp := httptest.NewRecorder()
	setKnownLeader(resp, true)
	header := resp.Header().Get("X-Consul-KnownLeader")
	if header != "true" {
		t.Fatalf("Bad: %v", header)
	}
	resp = httptest.NewRecorder()
	setKnownLeader(resp, false)
	header = resp.Header().Get("X-Consul-KnownLeader")
	if header != "false" {
		t.Fatalf("Bad: %v", header)
	}
}

func TestSetLastContact(t *testing.T) {
	resp := httptest.NewRecorder()
	setLastContact(resp, 123456*time.Microsecond)
	header := resp.Header().Get("X-Consul-LastContact")
	if header != "123" {
		t.Fatalf("Bad: %v", header)
	}
}

func TestSetMeta(t *testing.T) {
	meta := structs.QueryMeta{
		Index:       1000,
		KnownLeader: true,
		LastContact: 123456 * time.Microsecond,
	}
	resp := httptest.NewRecorder()
	setMeta(resp, &meta)
	header := resp.Header().Get("X-Consul-Index")
	if header != "1000" {
		t.Fatalf("Bad: %v", header)
	}
	header = resp.Header().Get("X-Consul-KnownLeader")
	if header != "true" {
		t.Fatalf("Bad: %v", header)
	}
	header = resp.Header().Get("X-Consul-LastContact")
	if header != "123" {
		t.Fatalf("Bad: %v", header)
	}
}

func TestHTTPAPIResponseHeaders(t *testing.T) {
	dir, srv := makeHTTPServer(t)
	srv.agent.config.HTTPAPIResponseHeaders = map[string]string{
		"Access-Control-Allow-Origin": "*",
		"X-XSS-Protection":            "1; mode=block",
	}

	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	resp := httptest.NewRecorder()

	handler := func(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
		return nil, nil
	}

	req, _ := http.NewRequest("GET", "/v1/agent/self", nil)
	srv.wrap(handler)(resp, req)

	origin := resp.Header().Get("Access-Control-Allow-Origin")
	if origin != "*" {
		t.Fatalf("bad Access-Control-Allow-Origin: expected %q, got %q", "*", origin)
	}

	xss := resp.Header().Get("X-XSS-Protection")
	if xss != "1; mode=block" {
		t.Fatalf("bad X-XSS-Protection header: expected %q, got %q", "1; mode=block", xss)
	}
}

func TestContentTypeIsJSON(t *testing.T) {
	dir, srv := makeHTTPServer(t)

	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	resp := httptest.NewRecorder()

	handler := func(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
		// stub out a DirEntry so that it will be encoded as JSON
		return &structs.DirEntry{Key: "key"}, nil
	}

	req, _ := http.NewRequest("GET", "/v1/kv/key", nil)
	srv.wrap(handler)(resp, req)

	contentType := resp.Header().Get("Content-Type")

	if contentType != "application/json" {
		t.Fatalf("Content-Type header was not 'application/json'")
	}
}

func TestPrettyPrint(t *testing.T) {
	testPrettyPrint("pretty=1", t)
}

func TestPrettyPrintBare(t *testing.T) {
	testPrettyPrint("pretty", t)
}

func testPrettyPrint(pretty string, t *testing.T) {
	dir, srv := makeHTTPServer(t)
	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	r := &structs.DirEntry{Key: "key"}

	resp := httptest.NewRecorder()
	handler := func(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
		return r, nil
	}

	urlStr := "/v1/kv/key?" + pretty
	req, _ := http.NewRequest("GET", urlStr, nil)
	srv.wrap(handler)(resp, req)

	expected, _ := json.MarshalIndent(r, "", "    ")
	actual, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if !bytes.Equal(expected, actual) {
		t.Fatalf("bad: %q", string(actual))
	}
}

func TestParseWait(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?wait=60s&index=1000", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseWait(resp, req, &b); d {
		t.Fatalf("unexpected done")
	}

	if b.MinQueryIndex != 1000 {
		t.Fatalf("Bad: %v", b)
	}
	if b.MaxQueryTime != 60*time.Second {
		t.Fatalf("Bad: %v", b)
	}
}

func TestParseWait_InvalidTime(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?wait=60foo&index=1000", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseWait(resp, req, &b); !d {
		t.Fatalf("expected done")
	}

	if resp.Code != 400 {
		t.Fatalf("bad code: %v", resp.Code)
	}
}

func TestParseWait_InvalidIndex(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?wait=60s&index=foo", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseWait(resp, req, &b); !d {
		t.Fatalf("expected done")
	}

	if resp.Code != 400 {
		t.Fatalf("bad code: %v", resp.Code)
	}
}

func TestParseConsistency(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?stale", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseConsistency(resp, req, &b); d {
		t.Fatalf("unexpected done")
	}

	if !b.AllowStale {
		t.Fatalf("Bad: %v", b)
	}
	if b.RequireConsistent {
		t.Fatalf("Bad: %v", b)
	}

	b = structs.QueryOptions{}
	req, err = http.NewRequest("GET",
		"/v1/catalog/nodes?consistent", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseConsistency(resp, req, &b); d {
		t.Fatalf("unexpected done")
	}

	if b.AllowStale {
		t.Fatalf("Bad: %v", b)
	}
	if !b.RequireConsistent {
		t.Fatalf("Bad: %v", b)
	}
}

func TestParseConsistency_Invalid(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?stale&consistent", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseConsistency(resp, req, &b); !d {
		t.Fatalf("expected done")
	}

	if resp.Code != 400 {
		t.Fatalf("bad code: %v", resp.Code)
	}
}

// assertIndex tests that X-Consul-Index is set and non-zero
func assertIndex(t *testing.T, resp *httptest.ResponseRecorder) {
	header := resp.Header().Get("X-Consul-Index")
	if header == "" || header == "0" {
		t.Fatalf("Bad: %v", header)
	}
}

// checkIndex is like assertIndex but returns an error
func checkIndex(resp *httptest.ResponseRecorder) error {
	header := resp.Header().Get("X-Consul-Index")
	if header == "" || header == "0" {
		return fmt.Errorf("Bad: %v", header)
	}
	return nil
}

// getIndex parses X-Consul-Index
func getIndex(t *testing.T, resp *httptest.ResponseRecorder) uint64 {
	header := resp.Header().Get("X-Consul-Index")
	if header == "" {
		t.Fatalf("Bad: %v", header)
	}
	val, err := strconv.Atoi(header)
	if err != nil {
		t.Fatalf("Bad: %v", header)
	}
	return uint64(val)
}

func httpTest(t *testing.T, f func(srv *HTTPServer)) {
	dir, srv := makeHTTPServer(t)
	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()
	testutil.WaitForLeader(t, srv.agent.RPC, "dc1")
	f(srv)
}
<commit_msg>agent: test socket file overwrite<commit_after>package agent

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"testing"
	"time"

	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/testutil"
)

func makeHTTPServer(t *testing.T) (string, *HTTPServer) {
	return makeHTTPServerWithConfig(t, nil)
}

func makeHTTPServerWithConfig(t *testing.T, cb func(c *Config)) (string, *HTTPServer) {
	conf := nextConfig()
	if cb != nil {
		cb(conf)
	}

	dir, agent := makeAgent(t, conf)
	uiDir := filepath.Join(dir, "ui")
	if err := os.Mkdir(uiDir, 755); err != nil {
		t.Fatalf("err: %v", err)
	}
	conf.UiDir = uiDir
	servers, err := NewHTTPServers(agent, conf, agent.logOutput)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if len(servers) == 0 {
		t.Fatalf(fmt.Sprintf("Failed to make HTTP server"))
	}
	return dir, servers[0]
}

func encodeReq(obj interface{}) io.ReadCloser {
	buf := bytes.NewBuffer(nil)
	enc := json.NewEncoder(buf)
	enc.Encode(obj)
	return ioutil.NopCloser(buf)
}

func TestHTTPServer_UnixSocket(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.SkipNow()
	}

	tempDir, err := ioutil.TempDir("", "consul")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(tempDir)
	socket := filepath.Join(tempDir, "test.sock")

	dir, srv := makeHTTPServerWithConfig(t, func(c *Config) {
		c.Addresses.HTTP = "unix://" + socket
	})
	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	// Ensure the socket was created
	if _, err := os.Stat(socket); err != nil {
		t.Fatalf("err: %s", err)
	}

	// Ensure we can get a response from the socket.
	path, _ := unixSocketAddr(srv.agent.config.Addresses.HTTP)
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(_, _ string) (net.Conn, error) {
				return net.Dial("unix", path)
			},
		},
	}

	// This URL doesn't look like it makes sense, but the scheme (http://) and
	// the host (127.0.0.1) are required by the HTTP client library. In reality
	// this will just use the custom dialer and talk to the socket.
	resp, err := client.Get("http://127.0.0.1/v1/agent/self")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer resp.Body.Close()

	if body, err := ioutil.ReadAll(resp.Body); err != nil || len(body) == 0 {
		t.Fatalf("bad: %s %v", body, err)
	}
}

func TestHTTPServer_UnixSocket_OverwriteFile(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.SkipNow()
	}

	tempDir, err := ioutil.TempDir("", "consul")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(tempDir)
	socket := filepath.Join(tempDir, "test.sock")

	// Create a regular file at the socket path
	if err := ioutil.WriteFile(socket, []byte("hello world"), 0644); err != nil {
		t.Fatalf("err: %s", err)
	}
	fi, err := os.Stat(socket)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if !fi.Mode().IsRegular() {
		t.Fatalf("not a regular file: %s", socket)
	}

	// Try to start the server with the same path anyways.
	dir, srv := makeHTTPServerWithConfig(t, func(c *Config) {
		c.Addresses.HTTP = "unix://" + socket
	})
	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	// Check if the socket overwrote the file
	fi, err = os.Stat(socket)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if fi.Mode().IsRegular() {
		t.Fatalf("should have socket file: %s", socket)
	}
}

func TestSetIndex(t *testing.T) {
	resp := httptest.NewRecorder()
	setIndex(resp, 1000)
	header := resp.Header().Get("X-Consul-Index")
	if header != "1000" {
		t.Fatalf("Bad: %v", header)
	}
	setIndex(resp, 2000)
	if v := resp.Header()["X-Consul-Index"]; len(v) != 1 {
		t.Fatalf("bad: %#v", v)
	}
}

func TestSetKnownLeader(t *testing.T) {
	resp := httptest.NewRecorder()
	setKnownLeader(resp, true)
	header := resp.Header().Get("X-Consul-KnownLeader")
	if header != "true" {
		t.Fatalf("Bad: %v", header)
	}
	resp = httptest.NewRecorder()
	setKnownLeader(resp, false)
	header = resp.Header().Get("X-Consul-KnownLeader")
	if header != "false" {
		t.Fatalf("Bad: %v", header)
	}
}

func TestSetLastContact(t *testing.T) {
	resp := httptest.NewRecorder()
	setLastContact(resp, 123456*time.Microsecond)
	header := resp.Header().Get("X-Consul-LastContact")
	if header != "123" {
		t.Fatalf("Bad: %v", header)
	}
}

func TestSetMeta(t *testing.T) {
	meta := structs.QueryMeta{
		Index:       1000,
		KnownLeader: true,
		LastContact: 123456 * time.Microsecond,
	}
	resp := httptest.NewRecorder()
	setMeta(resp, &meta)
	header := resp.Header().Get("X-Consul-Index")
	if header != "1000" {
		t.Fatalf("Bad: %v", header)
	}
	header = resp.Header().Get("X-Consul-KnownLeader")
	if header != "true" {
		t.Fatalf("Bad: %v", header)
	}
	header = resp.Header().Get("X-Consul-LastContact")
	if header != "123" {
		t.Fatalf("Bad: %v", header)
	}
}

func TestHTTPAPIResponseHeaders(t *testing.T) {
	dir, srv := makeHTTPServer(t)
	srv.agent.config.HTTPAPIResponseHeaders = map[string]string{
		"Access-Control-Allow-Origin": "*",
		"X-XSS-Protection":            "1; mode=block",
	}

	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	resp := httptest.NewRecorder()

	handler := func(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
		return nil, nil
	}

	req, _ := http.NewRequest("GET", "/v1/agent/self", nil)
	srv.wrap(handler)(resp, req)

	origin := resp.Header().Get("Access-Control-Allow-Origin")
	if origin != "*" {
		t.Fatalf("bad Access-Control-Allow-Origin: expected %q, got %q", "*", origin)
	}

	xss := resp.Header().Get("X-XSS-Protection")
	if xss != "1; mode=block" {
		t.Fatalf("bad X-XSS-Protection header: expected %q, got %q", "1; mode=block", xss)
	}
}

func TestContentTypeIsJSON(t *testing.T) {
	dir, srv := makeHTTPServer(t)

	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	resp := httptest.NewRecorder()

	handler := func(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
		// stub out a DirEntry so that it will be encoded as JSON
		return &structs.DirEntry{Key: "key"}, nil
	}

	req, _ := http.NewRequest("GET", "/v1/kv/key", nil)
	srv.wrap(handler)(resp, req)

	contentType := resp.Header().Get("Content-Type")

	if contentType != "application/json" {
		t.Fatalf("Content-Type header was not 'application/json'")
	}
}

func TestPrettyPrint(t *testing.T) {
	testPrettyPrint("pretty=1", t)
}

func TestPrettyPrintBare(t *testing.T) {
	testPrettyPrint("pretty", t)
}

func testPrettyPrint(pretty string, t *testing.T) {
	dir, srv := makeHTTPServer(t)
	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()

	r := &structs.DirEntry{Key: "key"}

	resp := httptest.NewRecorder()
	handler := func(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
		return r, nil
	}

	urlStr := "/v1/kv/key?" + pretty
	req, _ := http.NewRequest("GET", urlStr, nil)
	srv.wrap(handler)(resp, req)

	expected, _ := json.MarshalIndent(r, "", "    ")
	actual, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if !bytes.Equal(expected, actual) {
		t.Fatalf("bad: %q", string(actual))
	}
}

func TestParseWait(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?wait=60s&index=1000", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseWait(resp, req, &b); d {
		t.Fatalf("unexpected done")
	}

	if b.MinQueryIndex != 1000 {
		t.Fatalf("Bad: %v", b)
	}
	if b.MaxQueryTime != 60*time.Second {
		t.Fatalf("Bad: %v", b)
	}
}

func TestParseWait_InvalidTime(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?wait=60foo&index=1000", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseWait(resp, req, &b); !d {
		t.Fatalf("expected done")
	}

	if resp.Code != 400 {
		t.Fatalf("bad code: %v", resp.Code)
	}
}

func TestParseWait_InvalidIndex(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?wait=60s&index=foo", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseWait(resp, req, &b); !d {
		t.Fatalf("expected done")
	}

	if resp.Code != 400 {
		t.Fatalf("bad code: %v", resp.Code)
	}
}

func TestParseConsistency(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?stale", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseConsistency(resp, req, &b); d {
		t.Fatalf("unexpected done")
	}

	if !b.AllowStale {
		t.Fatalf("Bad: %v", b)
	}
	if b.RequireConsistent {
		t.Fatalf("Bad: %v", b)
	}

	b = structs.QueryOptions{}
	req, err = http.NewRequest("GET",
		"/v1/catalog/nodes?consistent", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseConsistency(resp, req, &b); d {
		t.Fatalf("unexpected done")
	}

	if b.AllowStale {
		t.Fatalf("Bad: %v", b)
	}
	if !b.RequireConsistent {
		t.Fatalf("Bad: %v", b)
	}
}

func TestParseConsistency_Invalid(t *testing.T) {
	resp := httptest.NewRecorder()
	var b structs.QueryOptions

	req, err := http.NewRequest("GET",
		"/v1/catalog/nodes?stale&consistent", nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if d := parseConsistency(resp, req, &b); !d {
		t.Fatalf("expected done")
	}

	if resp.Code != 400 {
		t.Fatalf("bad code: %v", resp.Code)
	}
}

// assertIndex tests that X-Consul-Index is set and non-zero
func assertIndex(t *testing.T, resp *httptest.ResponseRecorder) {
	header := resp.Header().Get("X-Consul-Index")
	if header == "" || header == "0" {
		t.Fatalf("Bad: %v", header)
	}
}

// checkIndex is like assertIndex but returns an error
func checkIndex(resp *httptest.ResponseRecorder) error {
	header := resp.Header().Get("X-Consul-Index")
	if header == "" || header == "0" {
		return fmt.Errorf("Bad: %v", header)
	}
	return nil
}

// getIndex parses X-Consul-Index
func getIndex(t *testing.T, resp *httptest.ResponseRecorder) uint64 {
	header := resp.Header().Get("X-Consul-Index")
	if header == "" {
		t.Fatalf("Bad: %v", header)
	}
	val, err := strconv.Atoi(header)
	if err != nil {
		t.Fatalf("Bad: %v", header)
	}
	return uint64(val)
}

func httpTest(t *testing.T, f func(srv *HTTPServer)) {
	dir, srv := makeHTTPServer(t)
	defer os.RemoveAll(dir)
	defer srv.Shutdown()
	defer srv.agent.Shutdown()
	testutil.WaitForLeader(t, srv.agent.RPC, "dc1")
	f(srv)
}
<|endoftext|>


<commit_before>package parser

import (
	"testing"

	"github.com/fatih/color"
)

var items = map[itemType]string{
	itemVersion:  "itemVersion",
	itemOperator: "itemOperator",
	itemSet:      "itemSet",
	itemRange:    "itemRange",
	itemAdvanced: "itemAdvanced",
	itemError:    "itemError",
	itemEOF:      "itemEOF",
}

type results []itemType

type lexerTestables struct {
	expected bool
	value    string
	result   results
}

var constraints = []*lexerTestables{
	{true, "1.0.0 || >=2.5.0 || 5.0.0 - 7.2.3",
		results{
			itemVersion,
			itemRange,
			itemOperator,
			itemVersion,
			itemRange,
			itemVersion,
			itemAdvanced,
			itemVersion,
		},
	},
	// Operators
	{true, "=2.3.2",
		results{itemOperator, itemVersion},
	},
	{true, "<=1.2.3",
		results{itemOperator, itemVersion},
	},
	{true, ">=1.2.3",
		results{itemOperator, itemVersion},
	},
	// Sets
	{true, "5.3.5 4.3.5",
		results{itemVersion, itemSet, itemVersion},
	},
	//Ranges
	{true, "5.3.5||4.3.5",
		results{itemVersion, itemRange, itemVersion},
	},
	{true, "5.3.5 ||4.3.5",
		results{itemVersion, itemRange, itemVersion},
	},
	{true, "5.3.5|| 4.3.5",
		results{itemVersion, itemRange, itemVersion},
	},
	// Tilde and Caret Ranges
	{false, "~ 1.2.3",
		results{itemAdvanced},
	},
	{true, "~1.2.3",
		results{itemAdvanced, itemVersion},
	},
	{true, "^4.5.2-alpha.1",
		results{itemAdvanced, itemVersion},
	},
	{false, ">= 1.2.3",
		results{},
	},
	// X-Ranges
	{true, "*",
		results{itemAdvanced},
	},
	{false, "**",
		results{},
	},
	{true, "1.0",
		results{itemAdvanced},
	},
	{true, "1.x",
		results{itemAdvanced},
	},
	{true, "*.x",
		results{itemAdvanced},
	},
	{false, "1.x+98uihuhyg",
		results{},
	},
	{true, "1.*.2",
		results{itemAdvanced},
	},
	{true, "1.*.2 || 1.x.4",
		results{itemAdvanced, itemRange, itemAdvanced},
	},
	{true, "1.*.2-beta",
		results{itemAdvanced},
	},
	{true, "*.1.2",
		results{itemAdvanced},
	},
	{false, "1x.2.*",
		results{},
	},
	{false, "1.x2.*",
		results{},
	},
	{false, "1...1",
		results{},
	},
	{false, "1.x.",
		results{},
	},

	// Assorted syntax errors
	{false, "1.2.3 >=",
		results{itemVersion, itemSet},
	},
	{false, "5.3.5 |1| 4.3.5",
		results{itemVersion},
	},
	{false, "5. 4.4",
		results{},
	},
	{false, "<1<1",
		results{itemOperator},
	},
	{false, "<1||",
		results{itemOperator, itemVersion},
	},
	{false, "M",
		results{},
	},
}

func init() {
	// Appends appropriate end token based on expected result.
	for _, c := range constraints {
		if c.expected {
			c.result = append(c.result, itemEOF)
		} else {
			c.result = append(c.result, itemError)
		}
	}
}

var cyan = color.New(color.FgCyan).SprintFunc()
var yellow = color.New(color.FgYellow).SprintFunc()

func TestLexer(t *testing.T) {
	for _, c := range constraints {
		_, ch := lex(c.value)
		result := true
		x := 0
		for i := range ch {

			result = (i.typ != itemError)

			if i.typ != c.result[x] {
				t.Logf("lex(%v) => %v(%v), want %v \n", cyan(c.value), items[i.typ], yellow(i), items[c.result[x]])
			}
			x++
		}
		if result != c.expected {
			t.Logf("lex(%v) => %t, want %t \n", cyan(c.value), result, c.expected)
		}
	}
}

// Poor implementation, just for initial testing.
func BenchmarkLexerComplex(b *testing.B) {
	const VERSION = "1.0.0 || >=2.5.0 || 5.0.0 - 7.2.3 || ~4.3.1 ^2.1.1"

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		_, ch := lex(VERSION)
		for {
			_, ok := <-ch
			if ok == false {
				//fmt.Printf("%v: '%v' \n", items[s.typ], s)
				//} else {
				break
			}
		}
	}
}

// Poor implementation, just for initial testing.
func BenchmarkLexerSimple(b *testing.B) {
	const VERSION = "1.0.0"

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		_, ch := lex(VERSION)
		for {
			_, ok := <-ch
			if ok == false {
				//fmt.Printf("%v: '%v' \n", items[s.typ], s)
				//} else {
				break
			}
		}
	}
}
<commit_msg>Fixed "out of range" bug.<commit_after>package parser

import (
	"testing"

	"github.com/fatih/color"
)

var items = map[itemType]string{
	itemVersion:  "itemVersion",
	itemOperator: "itemOperator",
	itemSet:      "itemSet",
	itemRange:    "itemRange",
	itemAdvanced: "itemAdvanced",
	itemError:    "itemError",
	itemEOF:      "itemEOF",
}

type results []itemType

type lexerTestables struct {
	expected bool
	value    string
	result   results
}

var constraints = []*lexerTestables{
	{true, "1.0.0 || >=2.5.0 || 5.0.0 - 7.2.3",
		results{
			itemVersion,
			itemRange,
			itemOperator,
			itemVersion,
			itemRange,
			itemVersion,
			itemAdvanced,
			itemVersion,
		},
	},
	// Operators
	{true, "=2.3.2",
		results{itemOperator, itemVersion},
	},
	{true, "<=1.2.3",
		results{itemOperator, itemVersion},
	},
	{true, ">=1.2.3",
		results{itemOperator, itemVersion},
	},
	// Sets
	{true, "5.3.5 4.3.5",
		results{itemVersion, itemSet, itemVersion},
	},
	//Ranges
	{true, "5.3.5||4.3.5",
		results{itemVersion, itemRange, itemVersion},
	},
	{true, "5.3.5 ||4.3.5",
		results{itemVersion, itemRange, itemVersion},
	},
	{true, "5.3.5|| 4.3.5",
		results{itemVersion, itemRange, itemVersion},
	},
	// Tilde and Caret Ranges
	{false, "~ 1.2.3",
		results{itemAdvanced},
	},
	{true, "~1.2.3",
		results{itemAdvanced, itemVersion},
	},
	{true, "^4.5.2-alpha.1",
		results{itemAdvanced, itemVersion},
	},
	{false, ">= 1.2.3",
		results{},
	},
	// X-Ranges
	{true, "*",
		results{itemAdvanced},
	},
	{false, "**",
		results{},
	},
	{true, "1.0",
		results{itemAdvanced},
	},
	{true, "1.x",
		results{itemAdvanced},
	},
	{true, "*.x",
		results{itemAdvanced},
	},
	{false, "1.x+98uihuhyg",
		results{},
	},
	{true, "1.*.2",
		results{itemAdvanced},
	},
	{true, "1.*.2 || 1.x.4",
		results{itemAdvanced, itemRange, itemAdvanced},
	},
	{true, "1.*.2-beta",
		results{itemAdvanced},
	},
	{true, "*.1.2",
		results{itemAdvanced},
	},
	{false, "1x.2.*",
		results{},
	},
	{false, "1.x2.*",
		results{},
	},
	{false, "1...1",
		results{},
	},
	{false, "1.x.",
		results{},
	},

	// Assorted syntax errors
	{false, "1.2.3 >=",
		results{itemVersion, itemSet},
	},
	{false, "5.3.5 |1| 4.3.5",
		results{itemVersion},
	},
	{false, "5. 4.4",
		results{},
	},
	{false, "<1<1",
		results{itemOperator},
	},
	{false, "<1||",
		results{itemOperator, itemVersion},
	},
	{false, "M",
		results{},
	},
}

func init() {
	// Appends appropriate end token based on expected result.
	for _, c := range constraints {
		if c.expected {
			c.result = append(c.result, itemEOF)
		} else {
			c.result = append(c.result, itemError)
		}
	}
}

var cyan = color.New(color.FgCyan).SprintFunc()
var yellow = color.New(color.FgYellow).SprintFunc()

func TestLexer(t *testing.T) {
	for _, c := range constraints {
		_, ch := lex(c.value)
		result := true
		x := 0
		for i := range ch {

			result = (i.typ != itemError)
			if len(c.result) > x {
				if i.typ != c.result[x] {
					t.Logf("lex(%v) => %v(%v), want %v \n", cyan(c.value), items[i.typ], yellow(i), items[c.result[x]])
				}
			} else {
				t.Logf("lex(%v) => %v(%v), want <nil>\n", cyan(c.value), items[i.typ], yellow(i))
			}
			x++
		}
		if result != c.expected {
			t.Logf("lex(%v) => %t, want %t \n", cyan(c.value), result, c.expected)
		}
	}
}

// Poor implementation, just for initial testing.
func BenchmarkLexerComplex(b *testing.B) {
	const VERSION = "1.0.0 || >=2.5.0 || 5.0.0 - 7.2.3 || ~4.3.1 ^2.1.1"

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		_, ch := lex(VERSION)
		for {
			_, ok := <-ch
			if ok == false {
				//fmt.Printf("%v: '%v' \n", items[s.typ], s)
				//} else {
				break
			}
		}
	}
}

// Poor implementation, just for initial testing.
func BenchmarkLexerSimple(b *testing.B) {
	const VERSION = "1.0.0"

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		_, ch := lex(VERSION)
		for {
			_, ok := <-ch
			if ok == false {
				//fmt.Printf("%v: '%v' \n", items[s.typ], s)
				//} else {
				break
			}
		}
	}
}
<|endoftext|>


<commit_before>package parser

import (
	"fmt"
	"testing"
)

var items = map[itemType]string{
	0: "itemVersion",
	1: "itemOperator",
	2: "itemSet",
	3: "itemRange",
	4: "itemError",
	5: "itemEOF",
}

var constraints = []string{
	"1.0 || >=2.5.0 || 5.0.0 - 7.2.3",
	"~1.2.3",
	"^4.5.2-alpha.1",
	"=2.3.2",
	"<=1.2.3",
	"5.3.5||4.3.5",
	"5.3.5 ||4.3.5",
	"5.3.5|| 4.3.5",
	"5.3.5 4.3.5",
	">=1.2.3",
	">= 1.2.3",
	"M",
}

// Just for debugging, not a real test. REMOVE THIS.
func TestLexer(t *testing.T) {
	for _, c := range constraints {
		_, ch := lex(c)
		for {
			s, ok := <-ch
			if ok != false {
				fmt.Printf("%v: '%v' \n", items[s.typ], s)
			} else {
				break
			}
		}
	}
}

// Poor implementation, just for initial testing.
func BenchmarkLexerComplex(b *testing.B) {
	const VERSION = "1.0.0 || >=2.5.0 || 5.0.0 - 7.2.3 || ~4.3.1 ^2.1.1"

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		_, ch := lex(VERSION)
		for {
			_, ok := <-ch
			if ok == false {
				//fmt.Printf("%v: '%v' \n", items[s.typ], s)
				//} else {
				break
			}
		}
	}
}

// Poor implementation, just for initial testing.
func BenchmarkLexerSimple(b *testing.B) {
	const VERSION = "1.0.0"

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		_, ch := lex(VERSION)
		for {
			_, ok := <-ch
			if ok == false {
				//fmt.Printf("%v: '%v' \n", items[s.typ], s)
				//} else {
				break
			}
		}
	}
}
<commit_msg>Added itemAdvance to error message map.<commit_after>package parser

import (
	"fmt"
	"testing"
)

var items = map[itemType]string{
	itemVersion:  "itemVersion",
	itemOperator: "itemOperator",
	itemSet:      "itemSet",
	itemRange:    "itemRange",
	itemAdvanced: "itemAdvanced",
	itemError:    "itemError",
	itemEOF:      "itemEOF",
}

var constraints = []string{
	"1.0 || >=2.5.0 || 5.0.0 - 7.2.3",
	"~1.2.3",
	"^4.5.2-alpha.1",
	"=2.3.2",
	"<=1.2.3",
	"5.3.5||4.3.5",
	"5.3.5 ||4.3.5",
	"5.3.5|| 4.3.5",
	"5.3.5 4.3.5",
	">=1.2.3",
	">= 1.2.3",
	"M",
}

// Just for debugging, not a real test. REMOVE THIS.
func TestLexer(t *testing.T) {
	for _, c := range constraints {
		_, ch := lex(c)
		for {
			s, ok := <-ch
			if ok != false {
				fmt.Printf("%v: '%v' \n", items[s.typ], s)
			} else {
				break
			}
		}
	}
}

// Poor implementation, just for initial testing.
func BenchmarkLexerComplex(b *testing.B) {
	const VERSION = "1.0.0 || >=2.5.0 || 5.0.0 - 7.2.3 || ~4.3.1 ^2.1.1"

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		_, ch := lex(VERSION)
		for {
			_, ok := <-ch
			if ok == false {
				//fmt.Printf("%v: '%v' \n", items[s.typ], s)
				//} else {
				break
			}
		}
	}
}

// Poor implementation, just for initial testing.
func BenchmarkLexerSimple(b *testing.B) {
	const VERSION = "1.0.0"

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		_, ch := lex(VERSION)
		for {
			_, ok := <-ch
			if ok == false {
				//fmt.Printf("%v: '%v' \n", items[s.typ], s)
				//} else {
				break
			}
		}
	}
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/tls"
	"fmt"
	"net"
)

// ListenerFactory is the factory function to create a listener.
type ListenerFactory func(map[string]string) (net.Listener, map[string]string, error)

// BuiltinListeners is the list of built-in listener types.
var BuiltinListeners = map[string]ListenerFactory{
	"tcp": tcpListenerFactory,
}

// NewListener creates a new listener of the given type with the given
// configuration. The type is looked up in the BuiltinListeners map.
func NewListener(t string, config map[string]string) (net.Listener, map[string]string, error) {
	f, ok := BuiltinListeners[t]
	if !ok {
		return nil, nil, fmt.Errorf("unknown listener type: %s", t)
	}

	return f(config)
}

func listenerWrapTLS(
	ln net.Listener,
	props map[string]string,
	config map[string]string) (net.Listener, map[string]string, error) {
	props["tls"] = "disabled"

	if v, ok := config["tls_disable"]; ok && v != "" {
		return ln, props, nil
	}

	certFile, ok := config["tls_cert_file"]
	if !ok {
		return nil, nil, fmt.Errorf("'tls_cert_file' must be set")
	}

	keyFile, ok := config["tls_key_file"]
	if !ok {
		return nil, nil, fmt.Errorf("'tls_key_file' must be set")
	}

	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return nil, nil, fmt.Errorf("error loading TLS cert: %s", err)
	}

	tlsConf := &tls.Config{}
	tlsConf.Certificates = []tls.Certificate{cert}
	tlsConf.NextProtos = []string{"http/1.1"}
	tlsConf.MinVersion = tls.VersionTLS12 // Minimum version is TLS 1.2
	tlsConf.ClientAuth = tls.RequestClientCert

	ln = tls.NewListener(ln, tlsConf)
	props["tls"] = "enabled"
	return ln, props, nil
}
<commit_msg>Allow specifying a TLS minimum version<commit_after>package server

import (
	"crypto/tls"
	"fmt"
	"net"
)

// ListenerFactory is the factory function to create a listener.
type ListenerFactory func(map[string]string) (net.Listener, map[string]string, error)

// BuiltinListeners is the list of built-in listener types.
var BuiltinListeners = map[string]ListenerFactory{
	"tcp": tcpListenerFactory,
}

// NewListener creates a new listener of the given type with the given
// configuration. The type is looked up in the BuiltinListeners map.
func NewListener(t string, config map[string]string) (net.Listener, map[string]string, error) {
	f, ok := BuiltinListeners[t]
	if !ok {
		return nil, nil, fmt.Errorf("unknown listener type: %s", t)
	}

	return f(config)
}

func listenerWrapTLS(
	ln net.Listener,
	props map[string]string,
	config map[string]string) (net.Listener, map[string]string, error) {
	props["tls"] = "disabled"

	if v, ok := config["tls_disable"]; ok && v != "" {
		return ln, props, nil
	}

	certFile, ok := config["tls_cert_file"]
	if !ok {
		return nil, nil, fmt.Errorf("'tls_cert_file' must be set")
	}

	keyFile, ok := config["tls_key_file"]
	if !ok {
		return nil, nil, fmt.Errorf("'tls_key_file' must be set")
	}

	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return nil, nil, fmt.Errorf("error loading TLS cert: %s", err)
	}

	tlslookup := map[string]uint16{
		"tls10": tls.VersionTLS10,
		"tls11": tls.VersionTLS11,
		"tls12": tls.VersionTLS12,
	}

	tlsvers, ok := config["tls_min_vers"]
	if !ok {
		tlsvers = "tls12"
	}

	tlsConf := &tls.Config{}
	tlsConf.Certificates = []tls.Certificate{cert}
	tlsConf.NextProtos = []string{"http/1.1"}
	tlsConf.MinVersion, ok = tlslookup[tlsvers]
	if !ok {
		return nil, nil, fmt.Errorf("'tls_min_vers' value %s not supported, please specify one of [tls10,tls11,tls12]", tlsvers)
	}
	tlsConf.ClientAuth = tls.RequestClientCert

	ln = tls.NewListener(ln, tlsConf)
	props["tls"] = "enabled"
	return ln, props, nil
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"fmt"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-log"
	"github.com/divisionone/go-micro/broker"
	"github.com/divisionone/go-micro/codec"
	"github.com/divisionone/go-micro/metadata"
	"github.com/divisionone/go-micro/registry"
	"github.com/divisionone/go-micro/transport"

	"github.com/divisionone/util/go/lib/addr"
)

type rpcServer struct {
	rpc  *server
	exit chan chan error

	sync.RWMutex
	opts        Options
	handlers    map[string]Handler
	subscribers map[*subscriber][]broker.Subscriber
	// used for first registration
	registered bool
	// graceful exit
	wg *sync.WaitGroup
}

func newRpcServer(opts ...Option) Server {
	options := newOptions(opts...)
	return &rpcServer{
		opts: options,
		rpc: &server{
			name:         options.Name,
			serviceMap:   make(map[string]*service),
			hdlrWrappers: options.HdlrWrappers,
		},
		handlers:    make(map[string]Handler),
		subscribers: make(map[*subscriber][]broker.Subscriber),
		exit:        make(chan chan error),
		wg:  		 wait(options.Context),
	}
}

func (s *rpcServer) accept(sock transport.Socket) {
	defer func() {
		// close socket
		sock.Close()

		if r := recover(); r != nil {
			log.Log("panic recovered: ", r)
			log.Log(string(debug.Stack()))
		}
	}()

	for {
		var msg transport.Message
		if err := sock.Recv(&msg); err != nil {
			return
		}

		// add to wait group if "wait" is opt-in
		if s.wg != nil {
			s.wg.Add(1)
		}

		// we use this Timeout header to set a server deadline
		to := msg.Header["Timeout"]
		// we use this Content-Type header to identify the codec needed
		ct := msg.Header["Content-Type"]

		cf, err := s.newCodec(ct)
		// TODO: needs better error handling
		if err != nil {
			sock.Send(&transport.Message{
				Header: map[string]string{
					"Content-Type": "text/plain",
				},
				Body: []byte(err.Error()),
			})
			if s.wg != nil {
				s.wg.Done()
			}
			return
		}

		codec := newRpcPlusCodec(&msg, sock, cf)

		// strip our headers
		hdr := make(map[string]string)
		for k, v := range msg.Header {
			hdr[k] = v
		}
		delete(hdr, "Content-Type")
		delete(hdr, "Timeout")

		ctx := metadata.NewContext(context.Background(), hdr)

		// set the timeout if we have it
		if len(to) > 0 {
			if n, err := strconv.ParseUint(to, 10, 64); err == nil {
				ctx, _ = context.WithTimeout(ctx, time.Duration(n))
			}
		}

		// TODO: needs better error handling
		if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
			if s.wg != nil {
				s.wg.Done()
			}
			return
		}

		if s.wg != nil {
			s.wg.Done()
		}
	}
}

func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
	if cf, ok := s.opts.Codecs[contentType]; ok {
		return cf, nil
	}
	if cf, ok := defaultCodecs[contentType]; ok {
		return cf, nil
	}
	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}

func (s *rpcServer) Options() Options {
	s.RLock()
	opts := s.opts
	s.RUnlock()
	return opts
}

func (s *rpcServer) Init(opts ...Option) error {
	s.Lock()
	for _, opt := range opts {
		opt(&s.opts)
	}
	// update internal server
	s.rpc = &server{
		name:         s.opts.Name,
		serviceMap:   s.rpc.serviceMap,
		hdlrWrappers: s.opts.HdlrWrappers,
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) NewHandler(h interface{}, opts ...HandlerOption) Handler {
	return newRpcHandler(h, opts...)
}

func (s *rpcServer) Handle(h Handler) error {
	s.Lock()
	defer s.Unlock()

	if err := s.rpc.register(h.Handler()); err != nil {
		return err
	}

	s.handlers[h.Name()] = h

	return nil
}

func (s *rpcServer) NewSubscriber(topic string, sb interface{}, opts ...SubscriberOption) Subscriber {
	return newSubscriber(topic, sb, opts...)
}

func (s *rpcServer) Subscribe(sb Subscriber) error {
	sub, ok := sb.(*subscriber)
	if !ok {
		return fmt.Errorf("invalid subscriber: expected *subscriber")
	}
	if len(sub.handlers) == 0 {
		return fmt.Errorf("invalid subscriber: no handler functions")
	}

	if err := validateSubscriber(sb); err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()
	_, ok = s.subscribers[sub]
	if ok {
		return fmt.Errorf("subscriber %v already exists", s)
	}
	s.subscribers[sub] = nil
	return nil
}

func (s *rpcServer) Register() error {
	// parse address for host, port
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
		partz := strings.Split(config.Address, ":")
		if len(partz) > 1 {
			port, _ = strconv.Atoi(partz[len(partz)-1])
		}
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	// register service
	node := &registry.Node{
		Id:       config.Name + "-" + config.Id,
		Address:  addr,
		Port:     port,
		Metadata: config.Metadata,
	}

	node.Metadata["transport"] = config.Transport.String()
	node.Metadata["broker"] = config.Broker.String()
	node.Metadata["server"] = s.String()
	node.Metadata["registry"] = config.Registry.String()

	s.RLock()
	// Maps are ordered randomly, sort the keys for consistency
	var handlerList []string
	for n, e := range s.handlers {
		// Only advertise non internal handlers
		if !e.Options().Internal {
			handlerList = append(handlerList, n)
		}
	}
	sort.Strings(handlerList)

	var subscriberList []*subscriber
	for e := range s.subscribers {
		// Only advertise non internal subscribers
		if !e.Options().Internal {
			subscriberList = append(subscriberList, e)
		}
	}
	sort.Slice(subscriberList, func(i, j int) bool {
		return subscriberList[i].topic > subscriberList[j].topic
	})

	var endpoints []*registry.Endpoint
	for _, n := range handlerList {
		endpoints = append(endpoints, s.handlers[n].Endpoints()...)
	}
	for _, e := range subscriberList {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	s.RUnlock()

	service := &registry.Service{
		Name:      config.Name,
		Version:   config.Version,
		Nodes:     []*registry.Node{node},
		Endpoints: endpoints,
	}

	s.Lock()
	registered := s.registered
	s.Unlock()

	if !registered {
		log.Logf("Registering node: %s", node.Id)
	}

	// create registry options
	rOpts := []registry.RegisterOption{registry.RegisterTTL(config.RegisterTTL)}

	if err := config.Registry.Register(service, rOpts...); err != nil {
		return err
	}

	// already registered? don't need to register subscribers
	if registered {
		return nil
	}

	s.Lock()
	defer s.Unlock()

	s.registered = true

	for sb, _ := range s.subscribers {
		handler := s.createSubHandler(sb, s.opts)
		var opts []broker.SubscribeOption
		if queue := sb.Options().Queue; len(queue) > 0 {
			opts = append(opts, broker.Queue(queue))
		}
		sub, err := config.Broker.Subscribe(sb.Topic(), handler, opts...)
		if err != nil {
			return err
		}
		s.subscribers[sb] = []broker.Subscriber{sub}
	}

	return nil
}

func (s *rpcServer) Deregister() error {
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	node := &registry.Node{
		Id:      config.Name + "-" + config.Id,
		Address: addr,
		Port:    port,
	}

	service := &registry.Service{
		Name:    config.Name,
		Version: config.Version,
		Nodes:   []*registry.Node{node},
	}

	log.Logf("Deregistering node: %s", node.Id)
	if err := config.Registry.Deregister(service); err != nil {
		return err
	}

	s.Lock()

	if !s.registered {
		s.Unlock()
		return nil
	}

	s.registered = false

	for sb, subs := range s.subscribers {
		for _, sub := range subs {
			log.Logf("Unsubscribing from topic: %s", sub.Topic())
			sub.Unsubscribe()
		}
		s.subscribers[sb] = nil
	}

	s.Unlock()
	return nil
}

func (s *rpcServer) Start() error {
	registerDebugHandler(s)
	config := s.Options()

	ts, err := config.Transport.Listen(config.Address)
	if err != nil {
		return err
	}

	log.Logf("Listening on %s", ts.Addr())
	s.Lock()
	s.opts.Address = ts.Addr()
	s.Unlock()

	go ts.Accept(s.accept)

	go func() {
		// wait for exit
		ch := <-s.exit

		// wait for requests to finish
		if s.wg != nil {
			s.wg.Wait()
		}

		// close transport listener
		ch <- ts.Close()

		// disconnect the broker
		config.Broker.Disconnect()
	}()

	// TODO: subscribe to cruft
	return config.Broker.Connect()
}

func (s *rpcServer) Stop() error {
	ch := make(chan error)
	s.exit <- ch
	return <-ch
}

func (s *rpcServer) String() string {
	return "rpc"
}
<commit_msg>gofmt<commit_after>package server

import (
	"context"
	"fmt"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-log"
	"github.com/divisionone/go-micro/broker"
	"github.com/divisionone/go-micro/codec"
	"github.com/divisionone/go-micro/metadata"
	"github.com/divisionone/go-micro/registry"
	"github.com/divisionone/go-micro/transport"

	"github.com/divisionone/util/go/lib/addr"
)

type rpcServer struct {
	rpc  *server
	exit chan chan error

	sync.RWMutex
	opts        Options
	handlers    map[string]Handler
	subscribers map[*subscriber][]broker.Subscriber
	// used for first registration
	registered bool
	// graceful exit
	wg *sync.WaitGroup
}

func newRpcServer(opts ...Option) Server {
	options := newOptions(opts...)
	return &rpcServer{
		opts: options,
		rpc: &server{
			name:         options.Name,
			serviceMap:   make(map[string]*service),
			hdlrWrappers: options.HdlrWrappers,
		},
		handlers:    make(map[string]Handler),
		subscribers: make(map[*subscriber][]broker.Subscriber),
		exit:        make(chan chan error),
		wg:          wait(options.Context),
	}
}

func (s *rpcServer) accept(sock transport.Socket) {
	defer func() {
		// close socket
		sock.Close()

		if r := recover(); r != nil {
			log.Log("panic recovered: ", r)
			log.Log(string(debug.Stack()))
		}
	}()

	for {
		var msg transport.Message
		if err := sock.Recv(&msg); err != nil {
			return
		}

		// add to wait group if "wait" is opt-in
		if s.wg != nil {
			s.wg.Add(1)
		}

		// we use this Timeout header to set a server deadline
		to := msg.Header["Timeout"]
		// we use this Content-Type header to identify the codec needed
		ct := msg.Header["Content-Type"]

		cf, err := s.newCodec(ct)
		// TODO: needs better error handling
		if err != nil {
			sock.Send(&transport.Message{
				Header: map[string]string{
					"Content-Type": "text/plain",
				},
				Body: []byte(err.Error()),
			})
			if s.wg != nil {
				s.wg.Done()
			}
			return
		}

		codec := newRpcPlusCodec(&msg, sock, cf)

		// strip our headers
		hdr := make(map[string]string)
		for k, v := range msg.Header {
			hdr[k] = v
		}
		delete(hdr, "Content-Type")
		delete(hdr, "Timeout")

		ctx := metadata.NewContext(context.Background(), hdr)

		// set the timeout if we have it
		if len(to) > 0 {
			if n, err := strconv.ParseUint(to, 10, 64); err == nil {
				ctx, _ = context.WithTimeout(ctx, time.Duration(n))
			}
		}

		// TODO: needs better error handling
		if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
			if s.wg != nil {
				s.wg.Done()
			}
			return
		}

		if s.wg != nil {
			s.wg.Done()
		}
	}
}

func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
	if cf, ok := s.opts.Codecs[contentType]; ok {
		return cf, nil
	}
	if cf, ok := defaultCodecs[contentType]; ok {
		return cf, nil
	}
	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}

func (s *rpcServer) Options() Options {
	s.RLock()
	opts := s.opts
	s.RUnlock()
	return opts
}

func (s *rpcServer) Init(opts ...Option) error {
	s.Lock()
	for _, opt := range opts {
		opt(&s.opts)
	}
	// update internal server
	s.rpc = &server{
		name:         s.opts.Name,
		serviceMap:   s.rpc.serviceMap,
		hdlrWrappers: s.opts.HdlrWrappers,
	}
	s.Unlock()
	return nil
}

func (s *rpcServer) NewHandler(h interface{}, opts ...HandlerOption) Handler {
	return newRpcHandler(h, opts...)
}

func (s *rpcServer) Handle(h Handler) error {
	s.Lock()
	defer s.Unlock()

	if err := s.rpc.register(h.Handler()); err != nil {
		return err
	}

	s.handlers[h.Name()] = h

	return nil
}

func (s *rpcServer) NewSubscriber(topic string, sb interface{}, opts ...SubscriberOption) Subscriber {
	return newSubscriber(topic, sb, opts...)
}

func (s *rpcServer) Subscribe(sb Subscriber) error {
	sub, ok := sb.(*subscriber)
	if !ok {
		return fmt.Errorf("invalid subscriber: expected *subscriber")
	}
	if len(sub.handlers) == 0 {
		return fmt.Errorf("invalid subscriber: no handler functions")
	}

	if err := validateSubscriber(sb); err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()
	_, ok = s.subscribers[sub]
	if ok {
		return fmt.Errorf("subscriber %v already exists", s)
	}
	s.subscribers[sub] = nil
	return nil
}

func (s *rpcServer) Register() error {
	// parse address for host, port
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
		partz := strings.Split(config.Address, ":")
		if len(partz) > 1 {
			port, _ = strconv.Atoi(partz[len(partz)-1])
		}
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	// register service
	node := &registry.Node{
		Id:       config.Name + "-" + config.Id,
		Address:  addr,
		Port:     port,
		Metadata: config.Metadata,
	}

	node.Metadata["transport"] = config.Transport.String()
	node.Metadata["broker"] = config.Broker.String()
	node.Metadata["server"] = s.String()
	node.Metadata["registry"] = config.Registry.String()

	s.RLock()
	// Maps are ordered randomly, sort the keys for consistency
	var handlerList []string
	for n, e := range s.handlers {
		// Only advertise non internal handlers
		if !e.Options().Internal {
			handlerList = append(handlerList, n)
		}
	}
	sort.Strings(handlerList)

	var subscriberList []*subscriber
	for e := range s.subscribers {
		// Only advertise non internal subscribers
		if !e.Options().Internal {
			subscriberList = append(subscriberList, e)
		}
	}
	sort.Slice(subscriberList, func(i, j int) bool {
		return subscriberList[i].topic > subscriberList[j].topic
	})

	var endpoints []*registry.Endpoint
	for _, n := range handlerList {
		endpoints = append(endpoints, s.handlers[n].Endpoints()...)
	}
	for _, e := range subscriberList {
		endpoints = append(endpoints, e.Endpoints()...)
	}
	s.RUnlock()

	service := &registry.Service{
		Name:      config.Name,
		Version:   config.Version,
		Nodes:     []*registry.Node{node},
		Endpoints: endpoints,
	}

	s.Lock()
	registered := s.registered
	s.Unlock()

	if !registered {
		log.Logf("Registering node: %s", node.Id)
	}

	// create registry options
	rOpts := []registry.RegisterOption{registry.RegisterTTL(config.RegisterTTL)}

	if err := config.Registry.Register(service, rOpts...); err != nil {
		return err
	}

	// already registered? don't need to register subscribers
	if registered {
		return nil
	}

	s.Lock()
	defer s.Unlock()

	s.registered = true

	for sb, _ := range s.subscribers {
		handler := s.createSubHandler(sb, s.opts)
		var opts []broker.SubscribeOption
		if queue := sb.Options().Queue; len(queue) > 0 {
			opts = append(opts, broker.Queue(queue))
		}
		sub, err := config.Broker.Subscribe(sb.Topic(), handler, opts...)
		if err != nil {
			return err
		}
		s.subscribers[sb] = []broker.Subscriber{sub}
	}

	return nil
}

func (s *rpcServer) Deregister() error {
	config := s.Options()
	var advt, host string
	var port int

	// check the advertise address first
	// if it exists then use it, otherwise
	// use the address
	if len(config.Advertise) > 0 {
		advt = config.Advertise
	} else {
		advt = config.Address
	}

	parts := strings.Split(advt, ":")
	if len(parts) > 1 {
		host = strings.Join(parts[:len(parts)-1], ":")
		port, _ = strconv.Atoi(parts[len(parts)-1])
	} else {
		host = parts[0]
	}

	addr, err := addr.Extract(host)
	if err != nil {
		return err
	}

	node := &registry.Node{
		Id:      config.Name + "-" + config.Id,
		Address: addr,
		Port:    port,
	}

	service := &registry.Service{
		Name:    config.Name,
		Version: config.Version,
		Nodes:   []*registry.Node{node},
	}

	log.Logf("Deregistering node: %s", node.Id)
	if err := config.Registry.Deregister(service); err != nil {
		return err
	}

	s.Lock()

	if !s.registered {
		s.Unlock()
		return nil
	}

	s.registered = false

	for sb, subs := range s.subscribers {
		for _, sub := range subs {
			log.Logf("Unsubscribing from topic: %s", sub.Topic())
			sub.Unsubscribe()
		}
		s.subscribers[sb] = nil
	}

	s.Unlock()
	return nil
}

func (s *rpcServer) Start() error {
	registerDebugHandler(s)
	config := s.Options()

	ts, err := config.Transport.Listen(config.Address)
	if err != nil {
		return err
	}

	log.Logf("Listening on %s", ts.Addr())
	s.Lock()
	s.opts.Address = ts.Addr()
	s.Unlock()

	go ts.Accept(s.accept)

	go func() {
		// wait for exit
		ch := <-s.exit

		// wait for requests to finish
		if s.wg != nil {
			s.wg.Wait()
		}

		// close transport listener
		ch <- ts.Close()

		// disconnect the broker
		config.Broker.Disconnect()
	}()

	// TODO: subscribe to cruft
	return config.Broker.Connect()
}

func (s *rpcServer) Stop() error {
	ch := make(chan error)
	s.exit <- ch
	return <-ch
}

func (s *rpcServer) String() string {
	return "rpc"
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/privacylab/talek/common"
	_ "github.com/privacylab/talek/pir/pircpu"
)

import "testing"

func fromEnvOrDefault(envKey string, defaultVal int) int {
	if os.Getenv(envKey) != "" {
		val, _ := strconv.Atoi(os.Getenv(envKey))
		return val
	}
	return defaultVal
}

func testConf() Config {
	return Config{
		Config: &common.Config{
			NumBuckets:         uint64(fromEnvOrDefault("NUM_BUCKETS", 512)),
			BucketDepth:        uint64(fromEnvOrDefault("BUCKET_DEPTH", 4)),
			DataSize:           uint64(fromEnvOrDefault("DATA_SIZE", 512)),
			BloomFalsePositive: 0.95,
			MaxLoadFactor:      0.95,
			LoadFactorStep:     0.02,
		},
		ReadBatch:        fromEnvOrDefault("BATCH_SIZE", 8),
		WriteInterval:    time.Second,
		ReadInterval:     time.Second,
		TrustDomain:      nil,
		TrustDomainIndex: 0,
	}
}

func TestShardSanity(t *testing.T) {
	shard := NewShard("Test Shard", "cpu.0", testConf())
	if shard == nil {
		t.Error("Failed to create shard.")
		return
	}

	writeReplyChan := make(chan *common.WriteReply)

	data := make([]byte, conf.Config.DataSize)
	copy(data, bytes.NewBufferString("Magic").Bytes())
	shard.Write(&common.ReplicaWriteArgs{
		WriteArgs: common.WriteArgs{
			Bucket1:        0,
			Bucket2:        1,
			Data:           data,
			InterestVector: []byte{},
			ReplyChan:      writeReplyChan,
		},
		EpochFlag: false,
	})

	// Force DB write.
	shard.syncChan <- 1

	replychan := make(chan *common.BatchReadReply)

	rv := make([]byte, 512)
	rv[0] = 0xff
	req := common.PirArgs{RequestVector: rv}
	reqs := make([]common.PirArgs, 8)
	for i := 0; i < 8; i++ {
		reqs[i] = req
	}
	shard.BatchRead(&DecodedBatchReadRequest{Args: reqs, ReplyChan: replychan})

	reply := <-replychan
	if reply.Replies[0].Data[0] != bytes.NewBufferString("Magic").Bytes()[0] {
		t.Error("Failed to round-trip a write.")
		return
	}

	shard.Close()
}

func BenchmarkShard(b *testing.B) {
	fmt.Printf("Benchmark began with N=%d\n", b.N)
	readsPerWrite := fromEnvOrDefault("READS_PER_WRITE", 20)

	conf := testConf()
	shard := NewShard("Test Shard", "cpu.0", conf)
	if shard == nil {
		b.Error("Failed to create shard.")
		return
	}

	replychan := make(chan *common.BatchReadReply)

	//A default write request
	stdWrite := common.WriteArgs{Bucket1: 0, Bucket2: 1, Data: bytes.NewBufferString("Magic").Bytes(), InterestVector: []byte{}}
	shardWrite := &common.ReplicaWriteArgs{
		WriteArgs: stdWrite,
		EpochFlag: false,
	}

	//A default read request
	reqs := make([]common.PirArgs, conf.ReadBatch)
	rv := make([]byte, int(conf.NumBuckets))
	for i := 0; i < len(rv); i++ {
		rv[i] = byte(rand.Int())
	}
	req := common.PirArgs{RequestVector: rv}
	for i := 0; i < conf.ReadBatch; i++ {
		reqs[i] = req
	}
	stdRead := &DecodedBatchReadRequest{reqs, replychan}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if i%readsPerWrite == 0 {
			stdWrite.Bucket1 = uint64(rand.Int()) % conf.NumBuckets
			stdWrite.Bucket2 = uint64(rand.Int()) % conf.NumBuckets
			shard.Write(shardWrite)
		} else {
			shard.BatchRead(stdRead)
			reply := <-replychan

			if reply == nil || reply.Err != "" {
				b.Error("Read failed.")
			}
		}
		b.SetBytes(int64(1))
	}

	fmt.Printf("Benchmark called close w N=%d\n", b.N)
	shard.Close()
}
<commit_msg>fix size issue in shard_test<commit_after>package server

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/privacylab/talek/common"
	_ "github.com/privacylab/talek/pir/pircpu"
)

import "testing"

func fromEnvOrDefault(envKey string, defaultVal int) int {
	if os.Getenv(envKey) != "" {
		val, _ := strconv.Atoi(os.Getenv(envKey))
		return val
	}
	return defaultVal
}

func testConf() Config {
	return Config{
		Config: &common.Config{
			NumBuckets:         uint64(fromEnvOrDefault("NUM_BUCKETS", 512)),
			BucketDepth:        uint64(fromEnvOrDefault("BUCKET_DEPTH", 4)),
			DataSize:           uint64(fromEnvOrDefault("DATA_SIZE", 512)),
			BloomFalsePositive: 0.95,
			MaxLoadFactor:      0.95,
			LoadFactorStep:     0.02,
		},
		ReadBatch:        fromEnvOrDefault("BATCH_SIZE", 8),
		WriteInterval:    time.Second,
		ReadInterval:     time.Second,
		TrustDomain:      nil,
		TrustDomainIndex: 0,
	}
}

func TestShardSanity(t *testing.T) {
	conf := testConf()
	shard := NewShard("Test Shard", "cpu.0", conf)
	if shard == nil {
		t.Error("Failed to create shard.")
		return
	}

	writeReplyChan := make(chan *common.WriteReply)

	data := make([]byte, conf.Config.DataSize)
	copy(data, bytes.NewBufferString("Magic").Bytes())
	shard.Write(&common.ReplicaWriteArgs{
		WriteArgs: common.WriteArgs{
			Bucket1:        0,
			Bucket2:        1,
			Data:           data,
			InterestVector: []byte{},
			ReplyChan:      writeReplyChan,
		},
		EpochFlag: false,
	})

	// Force DB write.
	shard.syncChan <- 1

	replychan := make(chan *common.BatchReadReply)

	rv := make([]byte, 512)
	rv[0] = 0xff
	req := common.PirArgs{RequestVector: rv}
	reqs := make([]common.PirArgs, 8)
	for i := 0; i < 8; i++ {
		reqs[i] = req
	}
	shard.BatchRead(&DecodedBatchReadRequest{Args: reqs, ReplyChan: replychan})

	reply := <-replychan
	if reply.Replies[0].Data[0] != bytes.NewBufferString("Magic").Bytes()[0] {
		t.Error("Failed to round-trip a write.")
		return
	}

	shard.Close()
}

func BenchmarkShard(b *testing.B) {
	fmt.Printf("Benchmark began with N=%d\n", b.N)
	readsPerWrite := fromEnvOrDefault("READS_PER_WRITE", 20)

	conf := testConf()
	shard := NewShard("Test Shard", "cpu.0", conf)
	if shard == nil {
		b.Error("Failed to create shard.")
		return
	}

	replychan := make(chan *common.BatchReadReply)

	//A default write request
	stdWrite := common.WriteArgs{Bucket1: 0, Bucket2: 1, Data: bytes.NewBufferString("Magic").Bytes(), InterestVector: []byte{}}
	shardWrite := &common.ReplicaWriteArgs{
		WriteArgs: stdWrite,
		EpochFlag: false,
	}

	//A default read request
	reqs := make([]common.PirArgs, conf.ReadBatch)
	rv := make([]byte, int(conf.NumBuckets))
	for i := 0; i < len(rv); i++ {
		rv[i] = byte(rand.Int())
	}
	req := common.PirArgs{RequestVector: rv}
	for i := 0; i < conf.ReadBatch; i++ {
		reqs[i] = req
	}
	stdRead := &DecodedBatchReadRequest{reqs, replychan}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if i%readsPerWrite == 0 {
			stdWrite.Bucket1 = uint64(rand.Int()) % conf.NumBuckets
			stdWrite.Bucket2 = uint64(rand.Int()) % conf.NumBuckets
			shard.Write(shardWrite)
		} else {
			shard.BatchRead(stdRead)
			reply := <-replychan

			if reply == nil || reply.Err != "" {
				b.Error("Read failed.")
			}
		}
		b.SetBytes(int64(1))
	}

	fmt.Printf("Benchmark called close w N=%d\n", b.N)
	shard.Close()
}
<|endoftext|>


<commit_before>package scheme

import (
	"fmt"
	"testing"
)

type tokenTypeTest struct {
	source string
	result rune
}

type tokenizeTest struct {
	source string
	result []string
}

var tokenTypeTests = []tokenTypeTest{
	{"100", IntToken},
	{"#f", BooleanToken},
	{"#t", BooleanToken},
	{"+", IdentifierToken},
	{"-", IdentifierToken},
	{"f2000", IdentifierToken},
	{"a0?!*/<=>:$%^&_~", IdentifierToken},
}

var tokenizeTests = []tokenizeTest{
	{"1", []string{"1"}},
	{"#f", []string{"#f"}},
	{"#t", []string{"#t"}},
	{"(+ 1)", []string{"(", "+", "1", ")"}},
	{"(+ 1 (+ 1))", []string{"(", "+", "1", "(", "+", "1", ")", ")"}},
	{"(+ (- 1)2)", []string{"(", "+", "(", "-", "1", ")", "2", ")"}},
	{"(* (/ 1)2)", []string{"(", "*", "(", "/", "1", ")", "2", ")"}},
	{"(number? 1)", []string{"(", "number?", "1", ")"}},
}

func TestTokenType(t *testing.T) {
	for _, test := range tokenTypeTests {
		l := NewLexer(test.source)

		actual := l.TokenType()
		if actual != test.result {
			t.Errorf("%s => %s; want %s", test.source, tokenTypeString(actual), tokenTypeString(test.result))
		}
	}
}

func TestNextToken(t *testing.T) {
	for _, test := range tokenizeTests {
		actual := tokenizedStrings(test.source)
		if !areTheSameStrings(actual, test.result) {
			t.Errorf("%s => %s; want %s", test.source, actual, test.result)
		}
	}
}

func tokenTypeString(tokenType rune) string {
	switch tokenType {
	case EOF:
		return "EOF"
	case IdentifierToken:
		return "IdentifierToken"
	case IntToken:
		return "IntToken"
	default:
		return fmt.Sprintf("%c", tokenType)
	}
}

func tokenizedStrings(source string) []string {
	l := NewLexer(source)
	tokens := []string{}

	for {
		token := l.NextToken()
		if token == "" {
			break
		}
		tokens = append(tokens, token)
	}

	return tokens
}

func areTheSameStrings(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}

	for i := 0; i < len(a); i++ {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}
<commit_msg>Refactor token tests<commit_after>package scheme

import (
	"fmt"
	"strings"
	"testing"
)

type tokenTypeTest struct {
	source string
	result rune
}

type tokenizeTest struct {
	source string
	result []string
}

var tokenTypeTests = []tokenTypeTest{
	{"100", IntToken},
	{"#f", BooleanToken},
	{"#t", BooleanToken},
	{"+", IdentifierToken},
	{"-", IdentifierToken},
	{"f2000", IdentifierToken},
	{"a0?!*/<=>:$%^&_~", IdentifierToken},
}

var tokenizeTests = []tokenizeTest{
	{"1", makeTokens("1")},
	{"#f", makeTokens("#f")},
	{"#t", makeTokens("#t")},
	{"(+ 1)", makeTokens("(,+,1,)")},
	{"(+ 1 (+ 1))", makeTokens("(,+,1,(,+,1,),)")},
	{"(+ (- 1)2)", makeTokens("(,+,(,-,1,),2,)")},
	{"(* (/ 1)2)", makeTokens("(,*,(,/,1,),2,)")},
	{"(number? 1)", makeTokens("(,number?,1,)")},
}

func TestTokenType(t *testing.T) {
	for _, test := range tokenTypeTests {
		l := NewLexer(test.source)

		actual := l.TokenType()
		if actual != test.result {
			t.Errorf("%s => %s; want %s", test.source, tokenTypeString(actual), tokenTypeString(test.result))
		}
	}
}

func TestNextToken(t *testing.T) {
	for _, test := range tokenizeTests {
		actual := tokenizedStrings(test.source)
		if !areTheSameStrings(actual, test.result) {
			t.Errorf("%s => %s; want %s", test.source, actual, test.result)
		}
	}
}

func tokenTypeString(tokenType rune) string {
	switch tokenType {
	case EOF:
		return "EOF"
	case IdentifierToken:
		return "IdentifierToken"
	case IntToken:
		return "IntToken"
	default:
		return fmt.Sprintf("%c", tokenType)
	}
}

func tokenizedStrings(source string) []string {
	l := NewLexer(source)
	tokens := []string{}

	for {
		token := l.NextToken()
		if token == "" {
			break
		}
		tokens = append(tokens, token)
	}

	return tokens
}

func areTheSameStrings(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}

	for i := 0; i < len(a); i++ {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func makeTokens(text string) []string {
	return strings.Split(text, ",")
}
<|endoftext|>


<commit_before>package protocol

import (
	"runtime"

	"v2ray.com/core/common/net"
	"v2ray.com/core/common/uuid"
)

type RequestCommand byte

const (
	RequestCommandTCP = RequestCommand(0x01)
	RequestCommandUDP = RequestCommand(0x02)
)

// RequestOption is the options of a request.
type RequestOption byte

const (
	// RequestOptionChunkStream indicates request payload is chunked. Each chunk consists of length, authentication and payload.
	RequestOptionChunkStream = RequestOption(0x01)

	// RequestOptionConnectionReuse indicates client side expects to reuse the connection.
	RequestOptionConnectionReuse = RequestOption(0x02)

	// RequestOptionCompressedStream indicates request payload is compressed.
	RequestOptionCompressedStream = RequestOption(0x04)
)

func (v RequestOption) Has(option RequestOption) bool {
	return (v & option) == option
}

func (v *RequestOption) Set(option RequestOption) {
	*v = (*v | option)
}

func (v *RequestOption) Clear(option RequestOption) {
	*v = (*v & (^option))
}

type Security byte

func (v Security) Is(t SecurityType) bool {
	return v == Security(t)
}

func NormSecurity(s Security) Security {
	if s.Is(SecurityType_UNKNOWN) {
		return Security(SecurityType_LEGACY)
	}
	return s
}

type RequestHeader struct {
	Version  byte
	Command  RequestCommand
	Option   RequestOption
	Security Security
	Port     net.Port
	Address  net.Address
	User     *User
}

func (v *RequestHeader) Destination() net.Destination {
	if v.Command == RequestCommandUDP {
		return net.UDPDestination(v.Address, v.Port)
	}
	return net.TCPDestination(v.Address, v.Port)
}

type ResponseOption byte

const (
	ResponseOptionConnectionReuse = ResponseOption(0x01)
)

func (v *ResponseOption) Set(option ResponseOption) {
	*v = (*v | option)
}

func (v ResponseOption) Has(option ResponseOption) bool {
	return (v & option) == option
}

func (v *ResponseOption) Clear(option ResponseOption) {
	*v = (*v & (^option))
}

type ResponseCommand interface{}

type ResponseHeader struct {
	Option  ResponseOption
	Command ResponseCommand
}

type CommandSwitchAccount struct {
	Host     net.Address
	Port     net.Port
	ID       *uuid.UUID
	AlterIds uint16
	Level    uint32
	ValidMin byte
}

func (v *SecurityConfig) AsSecurity() Security {
	if v == nil {
		return Security(SecurityType_LEGACY)
	}
	if v.Type == SecurityType_AUTO {
		if runtime.GOARCH == "amd64" || runtime.GOARCH == "s390x" {
			return Security(SecurityType_AES128_GCM)
		}
		return Security(SecurityType_CHACHA20_POLY1305)
	}
	return NormSecurity(Security(v.Type))
}
<commit_msg>session frame support<commit_after>package protocol

import (
	"runtime"

	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/uuid"
)

type RequestCommand byte

const (
	RequestCommandTCP = RequestCommand(0x01)
	RequestCommandUDP = RequestCommand(0x02)
)

// RequestOption is the options of a request.
type RequestOption byte

const (
	// RequestOptionChunkStream indicates request payload is chunked. Each chunk consists of length, authentication and payload.
	RequestOptionChunkStream = RequestOption(0x01)

	// RequestOptionConnectionReuse indicates client side expects to reuse the connection.
	RequestOptionConnectionReuse = RequestOption(0x02)

	// RequestOptionCompressedStream indicates request payload is compressed.
	RequestOptionCompressedStream = RequestOption(0x04)
)

func (v RequestOption) Has(option RequestOption) bool {
	return (v & option) == option
}

func (v *RequestOption) Set(option RequestOption) {
	*v = (*v | option)
}

func (v *RequestOption) Clear(option RequestOption) {
	*v = (*v & (^option))
}

type Security byte

func (v Security) Is(t SecurityType) bool {
	return v == Security(t)
}

func NormSecurity(s Security) Security {
	if s.Is(SecurityType_UNKNOWN) {
		return Security(SecurityType_LEGACY)
	}
	return s
}

type RequestHeader struct {
	Version  byte
	Command  RequestCommand
	Option   RequestOption
	Security Security
	Port     net.Port
	Address  net.Address
	User     *User
}

func (v *RequestHeader) Destination() net.Destination {
	if v.Command == RequestCommandUDP {
		return net.UDPDestination(v.Address, v.Port)
	}
	return net.TCPDestination(v.Address, v.Port)
}

type ResponseOption byte

const (
	ResponseOptionConnectionReuse = ResponseOption(0x01)
)

func (v *ResponseOption) Set(option ResponseOption) {
	*v = (*v | option)
}

func (v ResponseOption) Has(option ResponseOption) bool {
	return (v & option) == option
}

func (v *ResponseOption) Clear(option ResponseOption) {
	*v = (*v & (^option))
}

type ResponseCommand interface{}

type ResponseHeader struct {
	Option  ResponseOption
	Command ResponseCommand
}

type CommandSwitchAccount struct {
	Host     net.Address
	Port     net.Port
	ID       *uuid.UUID
	AlterIds uint16
	Level    uint32
	ValidMin byte
}

func (v *SecurityConfig) AsSecurity() Security {
	if v == nil {
		return Security(SecurityType_LEGACY)
	}
	if v.Type == SecurityType_AUTO {
		if runtime.GOARCH == "amd64" || runtime.GOARCH == "s390x" {
			return Security(SecurityType_AES128_GCM)
		}
		return Security(SecurityType_CHACHA20_POLY1305)
	}
	return NormSecurity(Security(v.Type))
}

type SessionFrameOption byte

const (
	SessionFrameOptionDefault  SessionFrameOption = 0x00
	SessionFrameOptionNew      SessionFrameOption = 0x01
	SessionFrameOptionContinue SessionFrameOption = 0x02
	SessionFrameOptionEnd      SessionFrameOption = 0x03
)

type SessionId uint32

const (
	DefaultSessionId SessionId = 0
)

type SessionFrame struct {
	Id      SessionId
	Option  SessionFrameOption
	Address net.Address
	Port    net.Port
	Payload *buf.Buffer
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package eks implements 'kubetest' deployer interface.
// It uses 'awstester' and 'kubectl' binaries, rather than importing
// EKS packages directly. This is to eliminate the need of dependency
// management, both in upstream and downstream.
package eks

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"time"

	"github.com/aws/awstester/eksconfig"
	"github.com/aws/awstester/eksdeployer"

	"k8s.io/test-infra/kubetest/process"
)

// deployer implements EKS deployer interface using "awstester" binary.
// Satisfies "k8s.io/test-infra/kubetest/main.go" 'deployer' and 'publisher" interfaces.
// Reference https://github.com/kubernetes/test-infra/blob/master/kubetest/main.go.
type deployer struct {
	stopc         chan struct{}
	cfg           *eksconfig.Config
	awsTesterPath string
	kubectlPath   string
	ctrl          *process.Control
}

// NewDeployer creates a new EKS deployer.
func NewDeployer(timeout time.Duration, verbose bool) (eksdeployer.Deployer, error) {
	cfg := eksconfig.NewDefault()
	err := cfg.UpdateFromEnvs()
	if err != nil {
		return nil, err
	}
	var f *os.File
	f, err = ioutil.TempFile(os.TempDir(), "awstester")
	if err != nil {
		return nil, err
	}
	outputPath := f.Name()
	f.Close()
	cfg.ConfigPath, err = filepath.Abs(outputPath)
	if err != nil {
		return nil, err
	}
	if err = cfg.Sync(); err != nil {
		return nil, err
	}

	dp := &deployer{
		stopc: make(chan struct{}),
		cfg:   cfg,
		ctrl: process.NewControl(
			timeout,
			time.NewTimer(timeout),
			time.NewTimer(timeout),
			verbose,
		),
	}

	dp.awsTesterPath, err = exec.LookPath("awstester")
	if err != nil {
		return nil, fmt.Errorf("cannot find 'awstester' executable (%v)", err)
	}
	dp.kubectlPath, err = exec.LookPath("kubectl")
	if err != nil {
		return nil, fmt.Errorf("cannot find 'kubectl' executable (%v)", err)
	}
	// TODO(gyuho): replace this kubernetes native Go client
	_, err = exec.LookPath("aws-iam-authenticator")
	if err != nil {
		return nil, fmt.Errorf("cannot find 'aws-iam-authenticator' executable (%v)", err)
	}

	return dp, nil
}

// Up creates a new EKS cluster.
func (dp *deployer) Up() (err error) {
	// "create cluster" command outputs cluster information
	// in the configuraion file (e.g. VPC ID, ALB DNS names, etc.)
	// this needs be reloaded for other deployer method calls
	createCmd := exec.Command(
		dp.awsTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"create",
		"cluster",
	)
	errc := make(chan error)
	go func() {
		_, oerr := dp.ctrl.Output(createCmd)
		errc <- oerr
	}()
	select {
	case <-dp.stopc:
		fmt.Fprintln(os.Stderr, "received stop signal, interrupting 'create cluster' command...")
		ierr := createCmd.Process.Signal(syscall.SIGINT)
		err = fmt.Errorf("'create cluster' command interrupted (interrupt error %v)", ierr)
	case err = <-errc:
	}
	return err
}

// Down tears down the existing EKS cluster.
func (dp *deployer) Down() (err error) {
	// reload configuration from disk to read the latest configuration
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"delete",
		"cluster",
	))
	return err
}

// IsUp returns an error if the cluster is not up and running.
func (dp *deployer) IsUp() (err error) {
	// reload configuration from disk to read the latest configuration
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"check",
		"cluster",
	))
	if err != nil {
		return err
	}
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	if dp.cfg.ClusterState.Status != "ACTIVE" {
		return fmt.Errorf("cluster %q status is %q",
			dp.cfg.ClusterName,
			dp.cfg.ClusterState.Status,
		)
	}
	return nil
}

// TestSetup checks if EKS testing cluster has been set up or not.
func (dp *deployer) TestSetup() error {
	return dp.IsUp()
}

// GetClusterCreated returns EKS cluster creation time and error (if any).
func (dp *deployer) GetClusterCreated(v string) (time.Time, error) {
	err := dp.IsUp()
	if err != nil {
		return time.Time{}, err
	}
	return dp.cfg.ClusterState.Created, nil
}

func (dp *deployer) GetWorkerNodeLogs() (err error) {
	// reload configuration from disk to read the latest configuration
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"test", "get-worker-node-logs",
	))
	return err
}

// DumpClusterLogs dumps all logs to artifact directory.
// Let default kubetest log dumper handle all artifact uploads.
// See https://github.com/kubernetes/test-infra/pull/9811/files#r225776067.
func (dp *deployer) DumpClusterLogs(artifactDir, _ string) (err error) {
	// reload configuration from disk to read the latest configuration
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"test", "get-worker-node-logs",
	))
	if err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"test", "dump-cluster-logs",
		artifactDir,
	))
	return err
}

// Stop stops ongoing operations.
// This is useful for local development.
// For example, one may run "Up" but have to cancel onging "Up"
// operation. Then, it can just send syscall.SIGINT to trigger "Stop".
func (dp *deployer) Stop() {
	close(dp.stopc)
}

// LoadConfig reloads configuration from disk to read the latest
// cluster configuration and its states.
func (dp *deployer) LoadConfig() (eksconfig.Config, error) {
	var err error
	dp.cfg, err = eksconfig.Load(dp.cfg.ConfigPath)
	return *dp.cfg, err
}
<commit_msg>kubetest/eks: download "aws-k8s-tester" on start<commit_after>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package eks implements 'kubetest' deployer interface.
// It uses 'aws-k8s-tester' and 'kubectl' binaries, rather than importing internal packages.
// All underlying implementation and external dependencies are compiled into one binary.
package eks

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"time"

	"github.com/aws/aws-k8s-tester/eksconfig"
	"github.com/aws/aws-k8s-tester/ekstester"

	"k8s.io/test-infra/kubetest/process"
	"k8s.io/test-infra/kubetest/util"
)

// deployer implements EKS deployer interface using "aws-k8s-tester" binary.
// Satisfies "k8s.io/test-infra/kubetest/main.go" 'deployer' and 'publisher" interfaces.
// Reference https://github.com/kubernetes/test-infra/blob/master/kubetest/main.go.
type deployer struct {
	stopc            chan struct{}
	cfg              *eksconfig.Config
	awsK8sTesterPath string
	kubectlPath      string
	ctrl             *process.Control
}

// NewDeployer creates a new EKS deployer.
func NewDeployer(timeout time.Duration, verbose bool) (ekstester.Deployer, error) {
	cfg := eksconfig.NewDefault()
	err := cfg.UpdateFromEnvs()
	if err != nil {
		return nil, err
	}
	var f *os.File
	f, err = ioutil.TempFile(os.TempDir(), "aws-k8s-tester")
	if err != nil {
		return nil, err
	}
	outputPath := f.Name()
	f.Close()
	cfg.ConfigPath, err = filepath.Abs(outputPath)
	if err != nil {
		return nil, err
	}
	if err = cfg.Sync(); err != nil {
		return nil, err
	}

	dp := &deployer{
		stopc: make(chan struct{}),
		cfg:   cfg,
		ctrl: process.NewControl(
			timeout,
			time.NewTimer(timeout),
			time.NewTimer(timeout),
			verbose,
		),
	}

	dp.awsK8sTesterPath, err = exec.LookPath("aws-k8s-tester")
	if err != nil {
		dp.awsK8sTesterPath = filepath.Join(os.TempDir(), "aws-k8s-tester")
		var f *os.File
		f, err = os.Create(dp.awsK8sTesterPath)
		if err != nil {
			return nil, fmt.Errorf("failed to create %q (%v)", dp.awsK8sTesterPath, err)
		}
		defer f.Close()
		if err = httpRead(cfg.AWSK8sTesterDownloadURL, f); err != nil {
			return nil, err
		}
		if err = util.EnsureExecutable(dp.awsK8sTesterPath); err != nil {
			return nil, err
		}
	}

	dp.kubectlPath, err = exec.LookPath("kubectl")
	if err != nil {
		return nil, fmt.Errorf("cannot find 'kubectl' executable (%v)", err)
	}

	// TODO(gyuho): replace this kubernetes native Go client
	_, err = exec.LookPath("aws-iam-authenticator")
	if err != nil {
		bin := filepath.Join(os.TempDir(), "aws-iam-authenticator")
		var f *os.File
		f, err = os.Create(bin)
		if err != nil {
			return nil, fmt.Errorf("failed to create %q (%v)", bin, err)
		}
		defer f.Close()
		if err = httpRead(cfg.AWSIAMAuthenticatorDownloadURL, f); err != nil {
			return nil, err
		}
		if err = util.EnsureExecutable(bin); err != nil {
			return nil, err
		}
		if err = os.Rename(bin, "/usr/local/bin/aws-iam-authenticator"); err != nil {
			return nil, err
		}
	}

	return dp, nil
}

// Up creates a new EKS cluster.
func (dp *deployer) Up() (err error) {
	// "create cluster" command outputs cluster information
	// in the configuraion file (e.g. VPC ID, ALB DNS names, etc.)
	// this needs be reloaded for other deployer method calls
	createCmd := exec.Command(
		dp.awsK8sTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"create",
		"cluster",
	)
	errc := make(chan error)
	go func() {
		_, oerr := dp.ctrl.Output(createCmd)
		errc <- oerr
	}()
	select {
	case <-dp.stopc:
		fmt.Fprintln(os.Stderr, "received stop signal, interrupting 'create cluster' command...")
		ierr := createCmd.Process.Signal(syscall.SIGINT)
		err = fmt.Errorf("'create cluster' command interrupted (interrupt error %v)", ierr)
	case err = <-errc:
	}
	return err
}

// Down tears down the existing EKS cluster.
func (dp *deployer) Down() (err error) {
	// reload configuration from disk to read the latest configuration
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsK8sTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"delete",
		"cluster",
	))
	return err
}

// IsUp returns an error if the cluster is not up and running.
func (dp *deployer) IsUp() (err error) {
	// reload configuration from disk to read the latest configuration
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsK8sTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"check",
		"cluster",
	))
	if err != nil {
		return err
	}
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	if dp.cfg.ClusterState.Status != "ACTIVE" {
		return fmt.Errorf("cluster %q status is %q",
			dp.cfg.ClusterName,
			dp.cfg.ClusterState.Status,
		)
	}
	return nil
}

// TestSetup checks if EKS testing cluster has been set up or not.
func (dp *deployer) TestSetup() error {
	return dp.IsUp()
}

// GetClusterCreated returns EKS cluster creation time and error (if any).
func (dp *deployer) GetClusterCreated(v string) (time.Time, error) {
	err := dp.IsUp()
	if err != nil {
		return time.Time{}, err
	}
	return dp.cfg.ClusterState.Created, nil
}

func (dp *deployer) GetWorkerNodeLogs() (err error) {
	// reload configuration from disk to read the latest configuration
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsK8sTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"test", "get-worker-node-logs",
	))
	return err
}

// DumpClusterLogs dumps all logs to artifact directory.
// Let default kubetest log dumper handle all artifact uploads.
// See https://github.com/kubernetes/test-infra/pull/9811/files#r225776067.
func (dp *deployer) DumpClusterLogs(artifactDir, _ string) (err error) {
	// reload configuration from disk to read the latest configuration
	if _, err = dp.LoadConfig(); err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsK8sTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"test", "get-worker-node-logs",
	))
	if err != nil {
		return err
	}
	_, err = dp.ctrl.Output(exec.Command(
		dp.awsK8sTesterPath,
		"eks",
		"--path="+dp.cfg.ConfigPath,
		"test", "dump-cluster-logs",
		artifactDir,
	))
	return err
}

// Stop stops ongoing operations.
// This is useful for local development.
// For example, one may run "Up" but have to cancel onging "Up"
// operation. Then, it can just send syscall.SIGINT to trigger "Stop".
func (dp *deployer) Stop() {
	close(dp.stopc)
}

// LoadConfig reloads configuration from disk to read the latest
// cluster configuration and its states.
func (dp *deployer) LoadConfig() (eksconfig.Config, error) {
	var err error
	dp.cfg, err = eksconfig.Load(dp.cfg.ConfigPath)
	return *dp.cfg, err
}

var httpTransport *http.Transport

func init() {
	httpTransport = new(http.Transport)
	httpTransport.RegisterProtocol("file", http.NewFileTransport(http.Dir("/")))
}

// curl -L [URL] | writer
func httpRead(u string, wr io.Writer) error {
	log.Printf("curl %s", u)
	cli := &http.Client{Transport: httpTransport}
	r, err := cli.Get(u)
	if err != nil {
		return err
	}
	defer r.Body.Close()
	if r.StatusCode >= 400 {
		return fmt.Errorf("%v returned %d", u, r.StatusCode)
	}
	_, err = io.Copy(wr, r.Body)
	return err
}
<|endoftext|>