feat!: refactor ws response writer to add validation logic (#33)

* refactor!: ws metrics response writer

* chore: remove unused interface.

* feat: add validate graphql request method

* chore: change varname

* chore: update .golangci.yml

* chore: returning error should be `io.EOF`

Author: vuongxuongminh

.air.toml

@@ -6,7 +6,7 @@ tmp_dir = "tmp"
   bin = "./tmp/gbox run -config ./Caddyfile.dist -watch"
   cmd = "go build -o ./tmp/gbox ./cmd"
   delay = 1000
-  exclude_dir = ["assets", "tmp", "vendor", "testdata", "internal/testserver"]
+  exclude_dir = ["assets", "tmp", "vendor", "testdata", "internal/testserver", "charts", "dist"]
   exclude_file = []
   exclude_regex = ["_test.go"]
   exclude_unchanged = false

.golangci.yml

@@ -32,6 +32,9 @@ linters:
     - wrapcheck
     - goerr113
     - gochecknoglobals
+    - execinquery
+    - exhaustruct
+    - nonamedreturns
 
     # deprecated
     - interfacer

handler_test.go

@@ -110,7 +110,7 @@ func (s *HandlerIntegrationTestSuite) TestIntrospection() {
 		"disabled": {
 			extraConfig:  "disabled_introspection true",
 			payload:      `{"query": "query { __schema { queryType { name } } }"}`,
-			expectedBody: `{"errors":[{"message":"introspection queries are not allowed"}]}`,
+			expectedBody: `{"errors":[{"message":"introspection query is not allowed"}]}`,
 		},
 	}
 

metrics.go

@@ -56,11 +56,6 @@ type Metrics struct {
 	cachingCount      *prometheus.CounterVec
 }
 
-type requestMetrics interface {
-	addMetricsBeginRequest(*graphql.Request)
-	addMetricsEndRequest(*graphql.Request, time.Duration)
-}
-
 type cachingMetrics interface {
 	addMetricsCaching(*graphql.Request, CachingStatus)
 }

router.go

@@ -25,6 +25,8 @@ const (
 	graphQLPath         = "/graphql"
 )
 
+var ErrNotAllowIntrospectionQuery = errors.New("introspection query is not allowed")
+
 func (h *Handler) initRouter() {
 	router := mux.NewRouter()
 	router.Path(graphQLPath).HeadersRegexp(
@@ -75,8 +77,8 @@ func (h *Handler) GraphQLOverWebsocketHandle(w http.ResponseWriter, r *http.Requ
 	}
 
 	n := r.Context().Value(nextHandlerCtxKey).(caddyhttp.Handler)
-	mr := newWebsocketMetricsResponseWriter(w, h.schema, h)
-	reporter.error = h.ReverseProxy.ServeHTTP(mr, r, n)
+	wsr := newWebsocketResponseWriter(w, h)
+	reporter.error = h.ReverseProxy.ServeHTTP(wsr, r, n)
 }
 
 // GraphQLHandle ensure GraphQL request is safe before forwarding to upstream and caching query result of it.
@@ -91,36 +93,23 @@ func (h *Handler) GraphQLHandle(w http.ResponseWriter, r *http.Request) {
 
 	gqlRequest, err := h.unmarshalHTTPRequest(r)
 	if err != nil {
-		h.logger.Debug("unmarshal GQL cachingRequest from http cachingRequest failure", zap.Error(err))
+		h.logger.Debug("can not unmarshal graphql request from http request", zap.Error(err))
 		reporter.error = writeResponseErrors(err, w)
 
 		return
 	}
 
-	isIntrospectQuery, _ := gqlRequest.IsIntrospectionQuery()
-	start := time.Now()
-
 	h.addMetricsBeginRequest(gqlRequest)
-	defer func() {
-		h.addMetricsEndRequest(gqlRequest, time.Since(start))
-	}()
+	defer func(startedAt time.Time) {
+		h.addMetricsEndRequest(gqlRequest, time.Since(startedAt))
+	}(time.Now())
 
-	if isIntrospectQuery && h.DisabledIntrospection {
-		reporter.error = writeResponseErrors(errors.New("introspection queries are not allowed"), w)
+	if err = h.validateGraphqlRequest(gqlRequest); err != nil {
+		reporter.error = writeResponseErrors(err, w)
 
 		return
 	}
 
-	if h.Complexity != nil {
-		requestErrors := h.Complexity.validateRequest(h.schema, gqlRequest)
-
-		if requestErrors.Count() > 0 {
-			reporter.error = writeResponseErrors(requestErrors, w)
-
-			return
-		}
-	}
-
 	n := r.Context().Value(nextHandlerCtxKey).(caddyhttp.Handler)
 
 	if h.Caching != nil {
@@ -161,6 +150,24 @@ func (h *Handler) unmarshalHTTPRequest(r *http.Request) (*graphql.Request, error
 	return gqlRequest, nil
 }
 
+func (h *Handler) validateGraphqlRequest(r *graphql.Request) error {
+	isIntrospectQuery, _ := r.IsIntrospectionQuery()
+
+	if isIntrospectQuery && h.DisabledIntrospection {
+		return ErrNotAllowIntrospectionQuery
+	}
+
+	if h.Complexity != nil {
+		requestErrors := h.Complexity.validateRequest(h.schema, r)
+
+		if requestErrors.Count() > 0 {
+			return requestErrors
+		}
+	}
+
+	return nil
+}
+
 // AdminGraphQLHandle purging query result cached and describe cache key.
 func (h *Handler) AdminGraphQLHandle(w http.ResponseWriter, r *http.Request) {
 	resolver := admin.NewResolver(h.schema, h.schemaDocument, h.logger, h.Caching)

ws.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/json"
+	"io"
 	"net"
 	"net/http"
 	"time"
@@ -13,54 +14,79 @@ import (
 	"github.com/jensneuse/graphql-go-tools/pkg/graphql"
 )
 
-type wsMetricsResponseWriter struct {
-	requestMetrics
+type wsSubscriber interface {
+	onWsSubscribe(*graphql.Request) error
+	onWsClose(*graphql.Request, time.Duration)
+}
+
+func (h *Handler) onWsSubscribe(r *graphql.Request) (err error) {
+	if err = normalizeGraphqlRequest(h.schema, r); err != nil {
+		return err
+	}
+
+	if err = h.validateGraphqlRequest(r); err != nil {
+		return err
+	}
+
+	h.addMetricsBeginRequest(r)
+
+	return nil
+}
+
+func (h *Handler) onWsClose(r *graphql.Request, d time.Duration) {
+	h.addMetricsEndRequest(r, d)
+}
+
+type wsResponseWriter struct {
 	*caddyhttp.ResponseWriterWrapper
-	schema *graphql.Schema
+	subscriber wsSubscriber
 }
 
-func newWebsocketMetricsResponseWriter(w http.ResponseWriter, s *graphql.Schema, rm requestMetrics) *wsMetricsResponseWriter {
-	return &wsMetricsResponseWriter{
+func newWebsocketResponseWriter(w http.ResponseWriter, s wsSubscriber) *wsResponseWriter {
+	return &wsResponseWriter{
 		ResponseWriterWrapper: &caddyhttp.ResponseWriterWrapper{
 			ResponseWriter: w,
 		},
-		schema:         s,
-		requestMetrics: rm,
+		subscriber: s,
 	}
 }
 
-// Hijack connection for collecting subscription metrics.
-func (r *wsMetricsResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+// Hijack connection for validating, and collecting metrics.
+func (r *wsResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	c, w, e := r.ResponseWriterWrapper.Hijack()
 
 	if c != nil {
-		c = &wsMetricsConn{
-			Conn:           c,
-			requestMetrics: r.requestMetrics,
-			schema:         r.schema,
+		c = &wsConn{
+			Conn:         c,
+			wsSubscriber: r.subscriber,
 		}
 	}
 
 	return c, w, e
 }
 
-type wsMetricsConn struct {
+type wsConn struct {
 	net.Conn
-	requestMetrics
+	wsSubscriber
 	request     *graphql.Request
-	schema      *graphql.Schema
 	subscribeAt time.Time
 }
 
-func (c *wsMetricsConn) Read(b []byte) (n int, err error) {
+type wsMessage struct {
+	ID      interface{}     `json:"id"`
+	Type    string          `json:"type"`
+	Payload json.RawMessage `json:"payload,omitempty"`
+}
+
+func (c *wsConn) Read(b []byte) (n int, err error) {
 	n, err = c.Conn.Read(b)
 
-	if c.request != nil || err != nil {
-		if err != nil {
-			c.addMetricsEndRequest(c.request, time.Since(c.subscribeAt))
-			c.request = nil
-		}
+	if c.request != nil && err != nil {
+		c.onWsClose(c.request, time.Since(c.subscribeAt))
+		c.request = nil
+	}
 
+	if c.request != nil || err != nil {
 		return n, err
 	}
 
@@ -76,10 +102,7 @@ func (c *wsMetricsConn) Read(b []byte) (n int, err error) {
 	}
 
 	decoder := json.NewDecoder(r)
-	msg := &struct {
-		Type    string          `json:"type"`
-		Payload json.RawMessage `json:"payload"`
-	}{}
+	msg := &wsMessage{}
 
 	if e := decoder.Decode(msg); e != nil {
 		return n, err
@@ -92,14 +115,50 @@ func (c *wsMetricsConn) Read(b []byte) (n int, err error) {
 			return n, err
 		}
 
-		if e := normalizeGraphqlRequest(c.schema, request); e != nil {
-			return n, err
+		if e := c.onWsSubscribe(request); e != nil {
+			c.writeErrorMessage(msg.ID, e)
+			c.writeCompleteMessage(msg.ID)
+
+			return n, io.EOF
 		}
 
 		c.request = request
 		c.subscribeAt = time.Now()
-		c.addMetricsBeginRequest(request)
 	}
 
 	return n, err
 }
+
+func (c *wsConn) writeErrorMessage(id interface{}, errMsg error) error {
+	errMsgRaw, errMsgErr := json.Marshal(graphql.RequestErrorsFromError(errMsg))
+
+	if errMsgErr != nil {
+		return errMsgErr
+	}
+
+	msg := &wsMessage{
+		ID:      id,
+		Type:    "error",
+		Payload: json.RawMessage(errMsgRaw),
+	}
+
+	payload, err := json.Marshal(msg)
+	if err != nil {
+		return err
+	}
+
+	return wsutil.WriteServerText(c, payload)
+}
+
+func (c *wsConn) writeCompleteMessage(id interface{}) error {
+	msg := &wsMessage{
+		ID:   id,
+		Type: "complete",
+	}
+	payload, err := json.Marshal(msg)
+	if err != nil {
+		return err
+	}
+
+	return wsutil.WriteServerText(c, payload)
+}