From 0a5ff8dfaf139de9ac1ac0cdaa7f797d1bcb05c1 Mon Sep 17 00:00:00 2001
From: Christine Dodrill <me@christine.website>
Date: Fri, 6 Oct 2017 10:54:09 -0700
Subject: [PATCH] internal/server: clean up NewServer

---
 internal/server/server.go      | 18 ++++++++----------
 internal/server/server_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 10 deletions(-)
 create mode 100644 internal/server/server_test.go

diff --git a/internal/server/server.go b/internal/server/server.go
index 4a8fb02..313e9c6 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -13,6 +13,7 @@ import (
 	proto "git.xeserv.us/xena/route/proto"
 	"github.com/Xe/ln"
 	"github.com/mtneug/pkg/ulid"
+	"github.com/oxtoacart/bpool"
 	kcp "github.com/xtaci/kcp-go"
 	"golang.org/x/crypto/acme/autocert"
 	"golang.org/x/net/context"
@@ -136,18 +137,15 @@ func New(cfg Config) (*Server, error) {
 		Manager: m,
 	}
 
-	go s.listenKCP(context.Background(), cfg.BackendKCPAddr, &tls.Config{
+	tc := &tls.Config{
 		GetCertificate: m.GetCertificate,
-	})
+	}
 
-	go s.listenTCP(context.Background(), cfg.BackendTCPAddr, &tls.Config{
-		GetCertificate: m.GetCertificate,
-	})
+	go s.listenKCP(context.Background(), cfg.BackendKCPAddr, tc)
+	go s.listenTCP(context.Background(), cfg.BackendTCPAddr, tc)
 
-	gs := grpc.NewServer(grpc.Creds(credentials.NewTLS(&tls.Config{
-		GetCertificate:     m.GetCertificate,
-		InsecureSkipVerify: true,
-	})))
+	// gRPC setup
+	gs := grpc.NewServer(grpc.Creds(credentials.NewTLS(tc)))
 
 	proto.RegisterBackendsServer(gs, &Backend{Server: s})
 	proto.RegisterRoutesServer(gs, &Route{Server: s})
@@ -194,7 +192,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		Director:      s.Director,
 		Transport:     s.ts,
 		FlushInterval: 1 * time.Second,
-		//BufferPool:    bpool.NewBytePool(256, 4096),
+		BufferPool:    bpool.NewBytePool(256, 4096),
 	}
 
 	rp.ServeHTTP(w, r)
diff --git a/internal/server/server_test.go b/internal/server/server_test.go
new file mode 100644
index 0000000..506fdb3
--- /dev/null
+++ b/internal/server/server_test.go
@@ -0,0 +1,29 @@
+package server
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestDirector(t *testing.T) {
+	s := &Server{}
+
+	req, err := http.NewRequest("GET", "https://cetacean.club/", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	req.Header.Add("X-Forwarded-For", "Rick-James")
+	req.Header.Add("X-Client-Ip", "56.32.51.84")
+
+	s.Director(req)
+
+	for _, header := range []string{"X-Forwarded-For", "X-Client-Ip"} {
+		t.Run(header, func(t *testing.T) {
+			val := req.Header.Get(header)
+			if val != "" {
+				t.Fatalf("expected header %q to have no value, got: %v", header, val)
+			}
+		})
+	}
+}