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

* fix: fixing audit logs for websocket connections

ee/audit/audit_ee.go

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

ee/audit/interceptor.go

@@ -23,6 +23,8 @@ import (
 	"net/http"
 
 	"google.golang.org/grpc"
+
+	"github.com/dgraph-io/dgraph/graphql/schema"
 )
 
 func AuditRequestGRPC(ctx context.Context, req interface{},
@@ -35,3 +37,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

@@ -17,8 +17,10 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/gorilla/websocket"
 	"io"
 	"io/ioutil"
+	"net"
 	"net/http"
 	"regexp"
 	"strconv"
@@ -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,40 @@ 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 user string
+	if token := req.Header.Get("X-Dgraph-AccessToken"); token != "" {
+		user = getUser(token, false)
+		namespace, _ = x.ExtractNamespaceFromJwt(token)
+	} 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

@@ -20,6 +20,7 @@ import (
 	"compress/gzip"
 	"context"
 	"encoding/json"
+	"github.com/dgraph-io/dgraph/ee/audit"
 	"io"
 	"io/ioutil"
 	"mime"
@@ -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

@@ -56,19 +56,22 @@ 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
 	}
-
 	namespace, ok := claims["namespace"].(float64)
 	if !ok {
 		return 0, errors.Errorf("namespace in claims is not valid:%v", namespace)
 	}
 	return uint64(namespace), nil
 }
+
+func ExtractJWTNamespace(ctx context.Context) (uint64, error) {
+	jwtString, err := ExtractJwt(ctx)
+	if err != nil {
+		return 0, err
+	}
+	return ExtractNamespaceFromJwt(jwtString)
+}