fix: fixing audit logs for websocket connections (#8048) (#8053) (#8572)

Fix audit logs for requests over websocket.
(cherry picked from commit 3f8cff9)

Co-authored-by: aman bansal <[email protected]>

edgraph/access_ee.go

@@ -1137,7 +1137,7 @@ func AuthGuardianOfTheGalaxy(ctx context.Context) error {
 	if !x.WorkerConfig.AclEnabled {
 		return nil
 	}
-	ns, err := x.ExtractJWTNamespace(ctx)
+	ns, err := x.ExtractNamespaceFrom(ctx)
 	if err != nil {
 		return errors.Wrap(err, "Authorize guardian of the galaxy, extracting jwt token, error:")
 	}

edgraph/server.go

@@ -1054,7 +1054,7 @@ func filterTablets(ctx context.Context, ms *pb.MembershipState) error {
 	if !x.WorkerConfig.AclEnabled {
 		return nil
 	}
-	namespace, err := x.ExtractJWTNamespace(ctx)
+	namespace, err := x.ExtractNamespaceFrom(ctx)
 	if err != nil {
 		return errors.Errorf("Namespace not found in JWT.")
 	}
@@ -1547,7 +1547,7 @@ func validateNamespace(ctx context.Context, tc *api.TxnContext) error {
 		return nil
 	}
 
-	ns, err := x.ExtractJWTNamespace(ctx)
+	ns, err := x.ExtractNamespaceFrom(ctx)
 	if err != nil {
 		return err
 	}
@@ -1571,7 +1571,7 @@ func (s *Server) CommitOrAbort(ctx context.Context, tc *api.TxnContext) (*api.Tx
 		return &api.TxnContext{}, errors.Errorf(
 			"StartTs cannot be zero while committing a transaction")
 	}
-	if ns, err := x.ExtractJWTNamespace(ctx); err == nil {
+	if ns, err := x.ExtractNamespaceFrom(ctx); err == nil {
 		annotateNamespace(span, ns)
 	}
 	annotateStartTs(span, tc.StartTs)

ee/audit/audit_ee.go

@@ -54,6 +54,7 @@ const (
 	PoorManAuth      = "PoorManAuth"
 	Grpc             = "Grpc"
 	Http             = "Http"
+	WebSocket        = "Websocket"
 )
 
 var auditor = &auditLogger{}

ee/audit/interceptor.go

@@ -2,7 +2,7 @@
 // +build oss
 
 /*
- * Copyright 2022 Dgraph Labs, Inc. and Contributors
+ * Copyright 2023 Dgraph Labs, Inc. and Contributors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,8 @@ import (
 	"context"
 	"net/http"
 
+	"github.com/dgraph-io/dgraph/graphql/schema"
+
 	"google.golang.org/grpc"
 )
 
@@ -36,3 +38,7 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 		next.ServeHTTP(w, r)
 	})
 }
+
+func AuditWebSockets(ctx context.Context, req *schema.Request) {
+	return
+}

ee/audit/interceptor_ee.go

@@ -2,7 +2,7 @@
 // +build !oss
 
 /*
- * Copyright 2022 Dgraph Labs, Inc. and Contributors
+ * Copyright 2023 Dgraph Labs, Inc. and Contributors
  *
  * Licensed under the Dgraph Community License (the "License"); you
  * may not use this file except in compliance with the License. You
@@ -20,23 +20,25 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"net"
 	"net/http"
 	"regexp"
 	"strconv"
 	"strings"
 	"sync/atomic"
 
+	"github.com/dgraph-io/dgraph/graphql/schema"
+	"github.com/dgraph-io/dgraph/x"
+	"github.com/dgraph-io/gqlparser/v2/ast"
+	"github.com/dgraph-io/gqlparser/v2/parser"
+
 	"github.com/golang/glog"
+	"github.com/gorilla/websocket"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
-
-	"github.com/dgraph-io/dgraph/graphql/schema"
-	"github.com/dgraph-io/dgraph/x"
-	"github.com/dgraph-io/gqlparser/v2/ast"
-	"github.com/dgraph-io/gqlparser/v2/parser"
 )
 
 const (
@@ -92,6 +94,19 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 			next.ServeHTTP(w, r)
 			return
 		}
+
+		// Websocket connection in graphQl happens differently. We only get access tokens and
+		// metadata in payload later once the connection is upgraded to correct protocol.
+		// Doc: https://github.com/apollographql/subscriptions-transport-ws/blob/v0.9.4/PROTOCOL.md
+		//
+		// Auditing for websocket connections will be handled by graphql/admin/http.go:154#Subscribe
+		for _, subprotocol := range websocket.Subprotocols(r) {
+			if subprotocol == "graphql-ws" {
+				next.ServeHTTP(w, r)
+				return
+			}
+		}
+
 		rw := NewResponseWriter(w)
 		var buf bytes.Buffer
 		tee := io.TeeReader(r.Body, &buf)
@@ -102,6 +117,45 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 	})
 }
 
+func AuditWebSockets(ctx context.Context, req *schema.Request) {
+	if atomic.LoadUint32(&auditEnabled) == 0 {
+		return
+	}
+
+	namespace := uint64(0)
+	var err error
+	var user string
+	// TODO(anurag): X-Dgraph-AccessToken should be exported as a constant
+	if token := req.Header.Get("X-Dgraph-AccessToken"); token != "" {
+		user = getUser(token, false)
+		namespace, err = x.ExtractNamespaceFromJwt(token)
+		if err != nil {
+			glog.Warningf("Error while auditing websockets: %s", err)
+		}
+	} else if token := req.Header.Get("X-Dgraph-AuthToken"); token != "" {
+		user = getUser(token, true)
+	} else {
+		user = getUser("", false)
+	}
+
+	ip := ""
+	if peerInfo, ok := peer.FromContext(ctx); ok {
+		ip, _, _ = net.SplitHostPort(peerInfo.Addr.String())
+	}
+
+	auditor.Audit(&AuditEvent{
+		User:        user,
+		Namespace:   namespace,
+		ServerHost:  x.WorkerConfig.MyAddr,
+		ClientHost:  ip,
+		Endpoint:    "/graphql",
+		ReqType:     WebSocket,
+		Req:         truncate(req.Query, maxReqLength),
+		Status:      http.StatusText(http.StatusOK),
+		QueryParams: nil,
+	})
+}
+
 func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo) {
 	clientHost := ""
 	if p, ok := peer.FromContext(ctx); ok {

graphql/admin/http.go

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Dgraph Labs, Inc. and Contributors
+ * Copyright 2023 Dgraph Labs, Inc. and Contributors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,6 +33,7 @@ import (
 	"go.opencensus.io/trace"
 
 	"github.com/dgraph-io/dgraph/edgraph"
+	"github.com/dgraph-io/dgraph/ee/audit"
 	"github.com/dgraph-io/dgraph/graphql/api"
 	"github.com/dgraph-io/dgraph/graphql/resolve"
 	"github.com/dgraph-io/dgraph/graphql/schema"
@@ -194,6 +195,8 @@ func (gs *graphqlSubscription) Subscribe(
 		Variables:     variableValues,
 		Header:        reqHeader,
 	}
+
+	audit.AuditWebSockets(ctx, req)
 	namespace := x.ExtractNamespaceHTTP(&http.Request{Header: reqHeader})
 	glog.Infof("namespace: %d. Got GraphQL request over websocket.", namespace)
 	// first load the schema, then do anything else

x/jwt_helper.go

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2022 Dgraph Labs, Inc. and Contributors
+ * Copyright 2017-2023 Dgraph Labs, Inc. and Contributors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@ package x
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/dgrijalva/jwt-go"
 	"github.com/pkg/errors"
@@ -56,19 +57,23 @@ func ExtractUserName(jwtToken string) (string, error) {
 	return userId, nil
 }
 
-func ExtractJWTNamespace(ctx context.Context) (uint64, error) {
-	jwtString, err := ExtractJwt(ctx)
-	if err != nil {
-		return 0, err
-	}
-	claims, err := ParseJWT(jwtString)
+func ExtractNamespaceFromJwt(jwtToken string) (uint64, error) {
+	claims, err := ParseJWT(jwtToken)
 	if err != nil {
-		return 0, err
+		return 0, errors.Wrap(err, "extracting namespace from JWT")
 	}
-
 	namespace, ok := claims["namespace"].(float64)
 	if !ok {
 		return 0, errors.Errorf("namespace in claims is not valid:%v", namespace)
 	}
 	return uint64(namespace), nil
 }
+
+func ExtractNamespaceFrom(ctx context.Context) (uint64, error) {
+	jwtString, err := ExtractJwt(ctx)
+
+	if err != nil {
+		return 0, fmt.Errorf("extracting namespace from JWT %w", err)
+	}
+	return ExtractNamespaceFromJwt(jwtString)
+}

x/x.go

@@ -260,7 +260,7 @@ func GqlErrorf(message string, args ...interface{}) *GqlError {
 func ExtractNamespaceHTTP(r *http.Request) uint64 {
 	ctx := AttachAccessJwt(context.Background(), r)
 	// Ignoring error because the default value is zero anyways.
-	namespace, _ := ExtractJWTNamespace(ctx)
+	namespace, _ := ExtractNamespaceFrom(ctx)
 	return namespace
 }
 
@@ -441,7 +441,7 @@ func AttachJWTNamespace(ctx context.Context) context.Context {
 		return AttachNamespace(ctx, GalaxyNamespace)
 	}
 
-	ns, err := ExtractJWTNamespace(ctx)
+	ns, err := ExtractNamespaceFrom(ctx)
 	if err == nil {
 		// Attach the namespace only if we got one from JWT.
 		// This preserves any namespace directly present in the context which is needed for
@@ -469,7 +469,7 @@ func AttachJWTNamespaceOutgoing(ctx context.Context) (context.Context, error) {
 	if !WorkerConfig.AclEnabled {
 		return AttachNamespaceOutgoing(ctx, GalaxyNamespace), nil
 	}
-	ns, err := ExtractJWTNamespace(ctx)
+	ns, err := ExtractNamespaceFrom(ctx)
 	if err != nil {
 		return ctx, err
 	}