Refactor logic, split such as can be re-used for posting data too

get_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"io/ioutil"
 	"net/http"
 	"os"
 	"time"
@@ -15,28 +14,6 @@ import (
 	. "github.com/onsi/gomega"
 )
 
-func writeRead(conn *websocket.Conn, input []byte) ([]byte, error) {
-	writer, err := conn.NextWriter(websocket.BinaryMessage)
-	if err != nil {
-		return nil, err
-	}
-
-	if _, err := writer.Write(input); err != nil {
-		return nil, err
-	}
-
-	if err := writer.Close(); err != nil {
-		return nil, err
-	}
-
-	_, reader, err := conn.NextReader()
-	if err != nil {
-		return nil, err
-	}
-
-	return ioutil.ReadAll(reader)
-}
-
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
@@ -53,30 +30,12 @@ func WSServer(ctx context.Context) {
 	m := http.NewServeMux()
 	m.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) {
 		conn, _ := upgrader.Upgrade(w, r, nil) // error ignored for sake of simplicity
-
 		for {
-
-			token := r.Header.Get("Authorization")
+			if err := AuthRequest(r, conn); err != nil {
+				fmt.Println("error", err.Error())
+				return
+			}
 			awesome := r.Header.Get("awesome")
-			ek, at, err := GetAttestationData(token)
-			if err != nil {
-				fmt.Println("error", err.Error())
-				return
-			}
-
-			secret, challenge, err := GenerateChallenge(ek, at)
-			if err != nil {
-				fmt.Println("error", err.Error())
-				return
-			}
-
-			resp, _ := writeRead(conn, challenge)
-
-			if err := ValidateChallenge(secret, resp); err != nil {
-				fmt.Println("error validating challenge", err.Error())
-				return
-			}
-
 			writer, _ := conn.NextWriter(websocket.BinaryMessage)
 			json.NewEncoder(writer).Encode(map[string]string{"foo": "bar", "header": awesome})
 		}
@@ -91,6 +50,66 @@ func WSServer(ctx context.Context) {
 	}()
 }
 
+// Mimics a WS server which accepts TPM Bearer token and receives data
+func WSServerReceiver(ctx context.Context, c chan map[string]string) {
+	s := http.Server{
+		Addr:         ":8080",
+		ReadTimeout:  10 * time.Second,
+		WriteTimeout: 10 * time.Second,
+	}
+
+	m := http.NewServeMux()
+	m.HandleFunc("/post", func(w http.ResponseWriter, r *http.Request) {
+		conn, _ := upgrader.Upgrade(w, r, nil) // error ignored for sake of simplicity
+		for {
+			if err := AuthRequest(r, conn); err != nil {
+				fmt.Println("error", err.Error())
+				return
+			}
+			defer conn.Close()
+
+			v := map[string]string{}
+			err := conn.ReadJSON(&v)
+			if err != nil {
+				fmt.Println("error", err.Error())
+				return
+			}
+			c <- v
+		}
+	})
+
+	s.Handler = m
+
+	go s.ListenAndServe()
+	go func() {
+		<-ctx.Done()
+		s.Shutdown(ctx)
+	}()
+}
+
+var _ = Describe("POST", func() {
+	Context("challenges", func() {
+		It("posts pubhash", func() {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			rec := make(chan map[string]string, 10)
+			WSServerReceiver(ctx, rec)
+
+			conn, err := Connection("http://localhost:8080/post", Emulated, WithSeed(1))
+			Expect(err).ToNot(HaveOccurred())
+
+			defer conn.Close()
+
+			err = conn.WriteJSON(map[string]string{"foo": "bar", "header": "foo"})
+			Expect(err).ToNot(HaveOccurred())
+
+			res := <-rec
+			Expect(res).To(Equal(map[string]string{"foo": "bar", "header": "foo"}))
+		})
+	})
+})
+
 var _ = Describe("GET", func() {
 	Context("challenges", func() {
 		It("fails for permissions", func() {