general: add unit tests for >80% coverage

Includes a few minor fixes:
- frontend: support setting port for WHOIS server
- proxy: fix handling of very long lines
- proxy: refactor IP allowlist logic, parse allow IP list at startup

proxy/bird.go

@@ -8,19 +8,23 @@ import (
 	"strings"
 )
 
+const MAX_LINE_SIZE = 1024
+
 // Read a line from bird socket, removing preceding status number, output it.
 // Returns if there are more lines.
 func birdReadln(bird io.Reader, w io.Writer) bool {
 	// Read from socket byte by byte, until reaching newline character
-	c := make([]byte, 1024, 1024)
+	c := make([]byte, MAX_LINE_SIZE)
 	pos := 0
 	for {
-		if pos >= 1024 {
+		// Leave one byte for newline character
+		if pos >= MAX_LINE_SIZE-1 {
 			break
 		}
 		_, err := bird.Read(c[pos : pos+1])
 		if err != nil {
-			panic(err)
+			w.Write([]byte(err.Error()))
+			return false
 		}
 		if c[pos] == byte('\n') {
 			break
@@ -29,6 +33,7 @@ func birdReadln(bird io.Reader, w io.Writer) bool {
 	}
 
 	c = c[:pos+1]
+	c[pos] = '\n'
 	// print(string(c[:]))
 
 	// Remove preceding status number, different situations

proxy/bird_test.go

@@ -0,0 +1,213 @@
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"path"
+	"strings"
+	"testing"
+
+	"github.com/magiconair/properties/assert"
+)
+
+type BirdServer struct {
+	t             *testing.T
+	expectedQuery string
+	response      string
+	server        net.Listener
+	socket        string
+	injectError   string
+}
+
+func (s *BirdServer) initSocket() {
+	tmpDir, err := ioutil.TempDir("", "bird-lgproxy-go-mock")
+	if err != nil {
+		s.t.Fatal(err)
+	}
+	s.socket = path.Join(tmpDir, "mock.socket")
+}
+
+func (s *BirdServer) Listen() {
+	s.initSocket()
+
+	var err error
+	s.server, err = net.Listen("unix", s.socket)
+	if err != nil {
+		s.t.Error(err)
+	}
+}
+
+func (s *BirdServer) Run() {
+	for {
+		conn, err := s.server.Accept()
+		if err != nil {
+			break
+		}
+		if conn == nil {
+			break
+		}
+
+		reader := bufio.NewReader(conn)
+
+		conn.Write([]byte("1234 Hello from mock bird\n"))
+
+		query, err := reader.ReadBytes('\n')
+		if err != nil {
+			break
+		}
+		if strings.TrimSpace(string(query)) != "restrict" {
+			s.t.Errorf("Did not restrict bird permissions")
+		}
+		if s.injectError == "restriction" {
+			conn.Write([]byte("1234 Restriction is disabled!\n"))
+		} else {
+			conn.Write([]byte("1234 Access restricted\n"))
+		}
+
+		query, err = reader.ReadBytes('\n')
+		if err != nil {
+			break
+		}
+		if strings.TrimSpace(string(query)) != s.expectedQuery {
+			s.t.Errorf("Query %s doesn't match expectation %s", string(query), s.expectedQuery)
+		}
+
+		responseList := strings.Split(s.response, "\n")
+		for i := range responseList {
+			if i == len(responseList)-1 {
+				if s.injectError == "eof" {
+					conn.Write([]byte("0000 " + responseList[i]))
+				} else {
+					conn.Write([]byte("0000 " + responseList[i] + "\n"))
+				}
+			} else {
+				conn.Write([]byte("1234 " + responseList[i] + "\n"))
+			}
+		}
+
+		conn.Close()
+	}
+}
+
+func (s *BirdServer) Close() {
+	if s.server == nil {
+		return
+	}
+	s.server.Close()
+}
+
+func TestBirdReadln(t *testing.T) {
+	input := strings.NewReader("1234 Bird Message\n")
+	var output bytes.Buffer
+	birdReadln(input, &output)
+
+	assert.Equal(t, output.String(), "Bird Message\n")
+}
+
+func TestBirdReadlnNoPrefix(t *testing.T) {
+	input := strings.NewReader(" Message without prefix\n")
+	var output bytes.Buffer
+	birdReadln(input, &output)
+
+	assert.Equal(t, output.String(), "Message without prefix\n")
+}
+
+func TestBirdReadlnVeryLongLine(t *testing.T) {
+	input := strings.NewReader(strings.Repeat("A", 4096))
+	var output bytes.Buffer
+	birdReadln(input, &output)
+
+	assert.Equal(t, output.String(), strings.Repeat("A", 1022)+"\n")
+}
+
+func TestBirdWriteln(t *testing.T) {
+	var output bytes.Buffer
+	birdWriteln(&output, "Test command")
+	assert.Equal(t, output.String(), "Test command\n")
+}
+
+func TestBirdHandlerWithoutQuery(t *testing.T) {
+	r := httptest.NewRequest(http.MethodGet, "/bird", nil)
+	w := httptest.NewRecorder()
+	birdHandler(w, r)
+}
+
+func TestBirdHandlerWithQuery(t *testing.T) {
+	server := BirdServer{
+		t:             t,
+		expectedQuery: "show protocols",
+		response:      "Mock Response\nSecond Line",
+		injectError:   "",
+	}
+
+	server.Listen()
+	go server.Run()
+	defer server.Close()
+
+	setting.birdSocket = server.socket
+
+	r := httptest.NewRequest(http.MethodGet, "/bird?q="+url.QueryEscape(server.expectedQuery), nil)
+	w := httptest.NewRecorder()
+	birdHandler(w, r)
+
+	assert.Equal(t, w.Code, http.StatusOK)
+	assert.Equal(t, w.Body.String(), server.response+"\n")
+}
+
+func TestBirdHandlerWithBadSocket(t *testing.T) {
+	setting.birdSocket = "/nonexistent.sock"
+
+	r := httptest.NewRequest(http.MethodGet, "/bird?q="+url.QueryEscape("mock"), nil)
+	w := httptest.NewRecorder()
+	birdHandler(w, r)
+
+	assert.Equal(t, w.Code, http.StatusInternalServerError)
+}
+
+func TestBirdHandlerWithoutRestriction(t *testing.T) {
+	server := BirdServer{
+		t:             t,
+		expectedQuery: "show protocols",
+		response:      "Mock Response",
+		injectError:   "restriction",
+	}
+
+	server.Listen()
+	go server.Run()
+	defer server.Close()
+
+	setting.birdSocket = server.socket
+
+	r := httptest.NewRequest(http.MethodGet, "/bird?q="+url.QueryEscape("mock"), nil)
+	w := httptest.NewRecorder()
+	birdHandler(w, r)
+
+	assert.Equal(t, w.Code, http.StatusInternalServerError)
+}
+
+func TestBirdHandlerEOF(t *testing.T) {
+	server := BirdServer{
+		t:             t,
+		expectedQuery: "show protocols",
+		response:      "Mock Response\nSecond Line",
+		injectError:   "eof",
+	}
+
+	server.Listen()
+	go server.Run()
+	defer server.Close()
+
+	setting.birdSocket = server.socket
+
+	r := httptest.NewRequest(http.MethodGet, "/bird?q="+url.QueryEscape("show protocols"), nil)
+	w := httptest.NewRecorder()
+	birdHandler(w, r)
+
+	assert.Equal(t, w.Code, http.StatusOK)
+	assert.Equal(t, w.Body.String(), "Mock Response\nEOF")
+}

proxy/main.go

@@ -21,39 +21,49 @@ func invalidHandler(httpW http.ResponseWriter, httpR *http.Request) {
 	httpW.Write([]byte("Invalid Request\n"))
 }
 
+func hasAccess(remoteAddr string) bool {
+	// setting.allowedIPs will always have at least one element because of how it's defined
+	if len(setting.allowedIPs) == 0 {
+		return true
+	}
+
+	if !strings.Contains(remoteAddr, ":") {
+		return false
+	}
+
+	// Remove port from IP and remove brackets that are around IPv6 addresses
+	remoteAddr = remoteAddr[0:strings.LastIndex(remoteAddr, ":")]
+	remoteAddr = strings.Trim(remoteAddr, "[]")
+
+	ipObject := net.ParseIP(remoteAddr)
+	if ipObject == nil {
+		return false
+	}
+
+	for _, allowedIP := range setting.allowedIPs {
+		if ipObject.Equal(allowedIP) {
+			return true
+		}
+	}
+
+	return false
+}
+
 // Access handler, check to see if client IP in allowed IPs, continue if it is, send to invalidHandler if not
 func accessHandler(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(httpW http.ResponseWriter, httpR *http.Request) {
-
-		// setting.allowedIPs will always have at least one element because of how it's defined
-		if setting.allowedIPs[0] == "" {
+		if hasAccess(httpR.RemoteAddr) {
 			next.ServeHTTP(httpW, httpR)
-			return
+		} else {
+			invalidHandler(httpW, httpR)
 		}
-
-		IPPort := httpR.RemoteAddr
-
-		// Remove port from IP and remove brackets that are around IPv6 addresses
-		requestIp := IPPort[0:strings.LastIndex(IPPort, ":")]
-		requestIp = strings.Replace(requestIp, "[", "", -1)
-		requestIp = strings.Replace(requestIp, "]", "", -1)
-
-		for _, allowedIP := range setting.allowedIPs {
-			if requestIp == allowedIP {
-				next.ServeHTTP(httpW, httpR)
-				return
-			}
-		}
-
-		invalidHandler(httpW, httpR)
-		return
 	})
 }
 
 type settingType struct {
 	birdSocket string
 	listen     string
-	allowedIPs []string
+	allowedIPs []net.IP
 	tr_bin     string
 	tr_flags   []string
 	tr_raw     bool

proxy/main_test.go

@@ -0,0 +1,78 @@
+package main
+
+import (
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/magiconair/properties/assert"
+)
+
+func TestHasAccessNotConfigured(t *testing.T) {
+	setting.allowedIPs = []net.IP{}
+	assert.Equal(t, hasAccess("whatever"), true)
+}
+
+func TestHasAccessAllowIPv4(t *testing.T) {
+	setting.allowedIPs = []net.IP{net.ParseIP("1.2.3.4")}
+	assert.Equal(t, hasAccess("1.2.3.4:4321"), true)
+}
+
+func TestHasAccessDenyIPv4(t *testing.T) {
+	setting.allowedIPs = []net.IP{net.ParseIP("4.3.2.1")}
+	assert.Equal(t, hasAccess("1.2.3.4:4321"), false)
+}
+
+func TestHasAccessAllowIPv6(t *testing.T) {
+	setting.allowedIPs = []net.IP{net.ParseIP("2001:db8::1")}
+	assert.Equal(t, hasAccess("[2001:db8::1]:4321"), true)
+}
+
+func TestHasAccessAllowIPv6DifferentForm(t *testing.T) {
+	setting.allowedIPs = []net.IP{net.ParseIP("2001:0db8::1")}
+	assert.Equal(t, hasAccess("[2001:db8::1]:4321"), true)
+}
+
+func TestHasAccessDenyIPv6(t *testing.T) {
+	setting.allowedIPs = []net.IP{net.ParseIP("2001:db8::2")}
+	assert.Equal(t, hasAccess("[2001:db8::1]:4321"), false)
+}
+
+func TestHasAccessBadClientIP(t *testing.T) {
+	setting.allowedIPs = []net.IP{net.ParseIP("1.2.3.4")}
+	assert.Equal(t, hasAccess("not an IP"), false)
+}
+
+func TestHasAccessBadClientIPPort(t *testing.T) {
+	setting.allowedIPs = []net.IP{net.ParseIP("1.2.3.4")}
+	assert.Equal(t, hasAccess("not an IP:not a port"), false)
+}
+
+func TestAccessHandlerAllow(t *testing.T) {
+	baseHandler := http.NotFoundHandler()
+	wrappedHandler := accessHandler(baseHandler)
+
+	r := httptest.NewRequest(http.MethodGet, "/mock", nil)
+	r.RemoteAddr = "1.2.3.4:4321"
+	w := httptest.NewRecorder()
+
+	setting.allowedIPs = []net.IP{net.ParseIP("1.2.3.4")}
+
+	wrappedHandler.ServeHTTP(w, r)
+	assert.Equal(t, w.Code, http.StatusNotFound)
+}
+
+func TestAccessHandlerDeny(t *testing.T) {
+	baseHandler := http.NotFoundHandler()
+	wrappedHandler := accessHandler(baseHandler)
+
+	r := httptest.NewRequest(http.MethodGet, "/mock", nil)
+	r.RemoteAddr = "1.2.3.4:4321"
+	w := httptest.NewRecorder()
+
+	setting.allowedIPs = []net.IP{net.ParseIP("4.3.2.1")}
+
+	wrappedHandler.ServeHTTP(w, r)
+	assert.Equal(t, w.Code, http.StatusInternalServerError)
+}

proxy/settings.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"net"
 	"strings"
 
 	"github.com/google/shlex"
@@ -66,9 +67,17 @@ func parseSettings() {
 	setting.listen = viperSettings.Listen
 
 	if viperSettings.AllowedIPs != "" {
-		setting.allowedIPs = strings.Split(viperSettings.AllowedIPs, ",")
+		for _, ip := range strings.Split(viperSettings.AllowedIPs, ",") {
+			ipObject := net.ParseIP(ip)
+			if ipObject == nil {
+				fmt.Printf("Parse IP %s failed\n", ip)
+				continue
+			}
+
+			setting.allowedIPs = append(setting.allowedIPs, ipObject)
+		}
 	} else {
-		setting.allowedIPs = []string{""}
+		setting.allowedIPs = []net.IP{}
 	}
 
 	var err error

proxy/settings_test.go

@@ -0,0 +1,8 @@
+package main
+
+import "testing"
+
+func TestParseSettings(t *testing.T) {
+	parseSettings()
+	// Good as long as it doesn't panic
+}

proxy/traceroute_test.go

@@ -0,0 +1,168 @@
+package main
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/magiconair/properties/assert"
+)
+
+func TestTracerouteArgsToString(t *testing.T) {
+	result := tracerouteArgsToString("traceroute", []string{
+		"-a",
+		"-b",
+		"-c",
+	}, []string{
+		"google.com",
+	})
+
+	assert.Equal(t, result, "traceroute -a -b -c google.com")
+}
+
+func TestTracerouteTryExecuteSuccess(t *testing.T) {
+	_, err := tracerouteTryExecute("sh", []string{
+		"-c",
+	}, []string{
+		"true",
+	})
+
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestTracerouteTryExecuteFail(t *testing.T) {
+	_, err := tracerouteTryExecute("sh", []string{
+		"-c",
+	}, []string{
+		"false",
+	})
+
+	if err == nil {
+		t.Error("Should trigger error, not triggered")
+	}
+}
+
+func TestTracerouteDetectSuccess(t *testing.T) {
+	result := tracerouteDetect("sh", []string{
+		"-c",
+		"true",
+	})
+
+	assert.Equal(t, result, true)
+}
+
+func TestTracerouteDetectFail(t *testing.T) {
+	result := tracerouteDetect("sh", []string{
+		"-c",
+		"false",
+	})
+
+	assert.Equal(t, result, false)
+}
+
+func TestTracerouteAutodetect(t *testing.T) {
+	pathBackup := os.Getenv("PATH")
+	os.Setenv("PATH", "")
+	defer os.Setenv("PATH", pathBackup)
+
+	setting.tr_bin = ""
+	setting.tr_flags = []string{}
+	tracerouteAutodetect()
+	// Should not panic
+}
+
+func TestTracerouteAutodetectExisting(t *testing.T) {
+	setting.tr_bin = "mock"
+	setting.tr_flags = []string{"mock"}
+	tracerouteAutodetect()
+	assert.Equal(t, setting.tr_bin, "mock")
+	assert.Equal(t, setting.tr_flags, []string{"mock"})
+}
+
+func TestTracerouteAutodetectFlagsOnly(t *testing.T) {
+	pathBackup := os.Getenv("PATH")
+	os.Setenv("PATH", "")
+	defer os.Setenv("PATH", pathBackup)
+
+	setting.tr_bin = "mock"
+	setting.tr_flags = nil
+	tracerouteAutodetect()
+
+	// Should not panic
+}
+
+func TestTracerouteHandlerWithoutQuery(t *testing.T) {
+	r := httptest.NewRequest(http.MethodGet, "/traceroute", nil)
+	w := httptest.NewRecorder()
+	tracerouteHandler(w, r)
+	assert.Equal(t, w.Code, http.StatusInternalServerError)
+	if !strings.Contains(w.Body.String(), "Invalid Request") {
+		t.Error("Did not get invalid request")
+	}
+}
+
+func TestTracerouteHandlerShlexError(t *testing.T) {
+	r := httptest.NewRequest(http.MethodGet, "/traceroute?q="+url.QueryEscape("\"1.1.1.1"), nil)
+	w := httptest.NewRecorder()
+	tracerouteHandler(w, r)
+	assert.Equal(t, w.Code, http.StatusInternalServerError)
+	if !strings.Contains(w.Body.String(), "parse") {
+		t.Error("Did not get parsing error message")
+	}
+}
+
+func TestTracerouteHandlerNoTracerouteFound(t *testing.T) {
+	setting.tr_bin = ""
+	setting.tr_flags = nil
+
+	r := httptest.NewRequest(http.MethodGet, "/traceroute?q="+url.QueryEscape("1.1.1.1"), nil)
+	w := httptest.NewRecorder()
+	tracerouteHandler(w, r)
+	assert.Equal(t, w.Code, http.StatusInternalServerError)
+	if !strings.Contains(w.Body.String(), "not supported") {
+		t.Error("Did not get not supported error message")
+	}
+}
+
+func TestTracerouteHandlerExecuteError(t *testing.T) {
+	setting.tr_bin = "sh"
+	setting.tr_flags = []string{"-c", "false"}
+	setting.tr_raw = true
+
+	r := httptest.NewRequest(http.MethodGet, "/traceroute?q="+url.QueryEscape("1.1.1.1"), nil)
+	w := httptest.NewRecorder()
+	tracerouteHandler(w, r)
+	assert.Equal(t, w.Code, http.StatusInternalServerError)
+	if !strings.Contains(w.Body.String(), "Error executing traceroute") {
+		t.Error("Did not get not execute error message")
+	}
+}
+
+func TestTracerouteHandlerRaw(t *testing.T) {
+	setting.tr_bin = "sh"
+	setting.tr_flags = []string{"-c", "echo Mock"}
+	setting.tr_raw = true
+
+	r := httptest.NewRequest(http.MethodGet, "/traceroute?q="+url.QueryEscape("1.1.1.1"), nil)
+	w := httptest.NewRecorder()
+	tracerouteHandler(w, r)
+	assert.Equal(t, w.Code, http.StatusOK)
+	assert.Equal(t, w.Body.String(), "Mock\n")
+}
+
+func TestTracerouteHandlerPostprocess(t *testing.T) {
+	setting.tr_bin = "sh"
+	setting.tr_flags = []string{"-c", "echo \"first line\n 2 *\nthird line\""}
+	setting.tr_raw = false
+
+	r := httptest.NewRequest(http.MethodGet, "/traceroute?q="+url.QueryEscape("1.1.1.1"), nil)
+	w := httptest.NewRecorder()
+	tracerouteHandler(w, r)
+	assert.Equal(t, w.Code, http.StatusOK)
+	assert.Equal(t, w.Body.String(), "first line\nthird line\n\n1 hops not responding.")
+}