fix: add ws protocol connection_terminate msg (#191)

* feat: add support for queires wo/space before '{' * fix: support subs-creds-in-vars with null vars * fix: add ws protocol connection_terminate msg
dosco · Mar 8, 2021 · d2f8f20 · d2f8f20
1 parent b1bd0f0
commit d2f8f20
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 14 deletions.
diff --git a/core/internal/qcode/utils.go b/core/internal/qcode/utils.go
@@ -26,7 +26,25 @@ func GetQType(gql string) (QType, string) {
 			continue
 
 		case b == '{':
-			switch tok {
+			t := tok
+			if s != -1 {
+				if t == "" {
+					t = gql[s:i]
+				} else {
+					n := gql[s:i]
+					if (bc % 2) == 0 {
+						switch t {
+						case "query":
+							return QTQuery, n
+						case "mutation":
+							return QTMutation, n
+						case "subscription":
+							return QTSubscription, n
+						}
+					}
+				}
+			}
+			switch t {
 			case "", "query":
 				return QTQuery, ""
 			case "mutation":

diff --git a/core/internal/qcode/utils_test.go b/core/internal/qcode/utils_test.go
@@ -1,6 +1,8 @@
 package qcode
 
-import "testing"
+import (
+	"testing"
+)
 
 func TestGetQType(t *testing.T) {
 	type args struct {
@@ -31,19 +33,29 @@ func TestGetQType(t *testing.T) {
 			args: args{gql: `fragment User on users { id name } query getStuff { query mutation(id: "query \"test1 '{") { id } subscription }`},
 			want: want{QTQuery, "getStuff"},
 		},
+		ts{
+			name: "fragment last, query with name",
+			args: args{gql: `query getStuff { query mutation(id: "query \"test1 '{") { id } subscription }fragment User on users { id name }`},
+			want: want{QTQuery, "getStuff"},
+		},
 		ts{
 			name: "mutation",
 			args: args{gql: `mutation { query mutation(id: "query {") { id } subscription }`},
 			want: want{QTMutation, ""},
 		},
+		ts{
+			name: "subscription",
+			args: args{gql: `subscription { query mutation(id: "query {") { id } subscription }`},
+			want: want{QTSubscription, ""},
+		},
 		ts{
 			name: "default query",
 			args: args{gql: ` { query mutation(id: "query {") { id } subscription }`},
 			want: want{QTQuery, ""},
 		},
 		ts{
 			name: "default query with comment",
-			args: args{gql: `# mutation is good 
+			args: args{gql: `# mutation is good
 				query { query mutation(id: "query") { id } subscription }`},
 			want: want{QTQuery, ""},
 		},
@@ -52,6 +64,53 @@ func TestGetQType(t *testing.T) {
 			args: args{gql: `# query is good query { query mutation(id: "query {{") { id } subscription }`},
 			want: want{QTUnknown, ""},
 		},
+		// tests without space after the op type
+		ts{
+			name: "query without space",
+			args: args{gql: `query{ query mutation(id: "query {") { id } subscription }`},
+			want: want{QTQuery, ""},
+		},
+		ts{
+			name: "query with name, without space",
+			args: args{gql: `query getStuff{ query mutation(id: "query \"test1 '{") { id } subscription }`},
+			want: want{QTQuery, "getStuff"},
+		},
+		ts{
+			name: "fragment first, query with name, without space",
+			args: args{gql: `fragment User on users { id name } query getStuff{ query mutation(id: "query \"test1 '{") { id } subscription }`},
+			want: want{QTQuery, "getStuff"},
+		},
+		ts{
+			name: "fragment last, query with name, without space",
+			args: args{gql: `query getStuff{ query mutation(id: "query \"test1 '{") { id } subscription }fragment User on users { id name }`},
+			want: want{QTQuery, "getStuff"},
+		},
+		ts{
+			name: "mutation without space",
+			args: args{gql: `mutation{ query mutation(id: "query {") { id } subscription }`},
+			want: want{QTMutation, ""},
+		},
+		ts{
+			name: "subscription without space",
+			args: args{gql: `subscription{ query mutation(id: "query {") { id } subscription }`},
+			want: want{QTSubscription, ""},
+		},
+		ts{
+			name: "default query without space",
+			args: args{gql: `{ query mutation(id: "query {") { id } subscription }`},
+			want: want{QTQuery, ""},
+		},
+		ts{
+			name: "default query with comment without space",
+			args: args{gql: `# mutation is good
+				query{ query mutation(id: "query") { id } subscription }`},
+			want: want{QTQuery, ""},
+		},
+		ts{
+			name: "failed query with comment, without space",
+			args: args{gql: `# query is good query{ query mutation(id: "query {{") { id } subscription }`},
+			want: want{QTUnknown, ""},
+		},
 	}
 
 	for _, tt := range tests {

diff --git a/internal/serv/ws.go b/internal/serv/ws.go
@@ -139,17 +139,16 @@ func (sc *ServConfig) apiV1Ws(w http.ResponseWriter, r *http.Request) {
 					UserID         interface{} `json:"X-User-ID"`
 				}
 				var x authHeaders
-				if err = json.Unmarshal(msg.Payload.Vars, &x); err != nil {
-					break
-				}
-				if x.UserIDProvider != "" {
-					ctx = context.WithValue(ctx, core.UserIDProviderKey, x.UserIDProvider)
-				}
-				if x.UserRole != "" {
-					ctx = context.WithValue(ctx, core.UserRoleKey, x.UserRole)
-				}
-				if x.UserID != nil {
-					ctx = context.WithValue(ctx, core.UserIDKey, x.UserID)
+				if err = json.Unmarshal(msg.Payload.Vars, &x); err == nil {
+					if x.UserIDProvider != "" {
+						ctx = context.WithValue(ctx, core.UserIDProviderKey, x.UserIDProvider)
+					}
+					if x.UserRole != "" {
+						ctx = context.WithValue(ctx, core.UserRoleKey, x.UserRole)
+					}
+					if x.UserID != nil {
+						ctx = context.WithValue(ctx, core.UserIDKey, x.UserID)
+					}
 				}
 			}
 
@@ -164,6 +163,11 @@ func (sc *ServConfig) apiV1Ws(w http.ResponseWriter, r *http.Request) {
 			done <- true
 			run = false
 
+		case "connection_terminate":
+			m.Unsubscribe()
+			done <- true
+			return
+
 		default:
 			fields := []zapcore.Field{
 				zap.String("msg_type", msg.Type),