From 5a771e69c9e4d2e6967dc1a98486e882231c9430 Mon Sep 17 00:00:00 2001
From: aman bansal <bansalaman2905@gmail.com>
Date: Fri, 24 Sep 2021 15:30:43 +0530
Subject: [PATCH] fix: fixing audit logs for websocket connections (#8048)
 (#8053)

* fix: fixing audit logs for websocket connections

(cherry picked from commit 3f8cff957b0ce34673f70ffd8793d0353bf2a910)
---
 ee/audit/audit_ee.go       |  1 +
 ee/audit/interceptor.go    |  6 +++++
 ee/audit/interceptor_ee.go | 49 ++++++++++++++++++++++++++++++++++++++
 graphql/admin/http.go      |  3 +++
 x/jwt_helper.go            | 17 +++++++------
 5 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index 837d75d7c54..3074098d844 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -53,6 +53,7 @@ const (
 	PoorManAuth      = "PoorManAuth"
 	Grpc             = "Grpc"
 	Http             = "Http"
+	WebSocket        = "Websocket"
 )
 
 var auditor = &auditLogger{}
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 1a8e69f5d93..45b76270003 100644
--- a/ee/audit/interceptor.go
+++ b/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
+}
diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
index e9de9a1b777..aa26d9027b6 100644
--- a/ee/audit/interceptor_ee.go
+++ b/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 {
diff --git a/graphql/admin/http.go b/graphql/admin/http.go
index 56fdf51e0cc..bc39751183a 100644
--- a/graphql/admin/http.go
+++ b/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"
@@ -193,6 +194,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
diff --git a/x/jwt_helper.go b/x/jwt_helper.go
index 9830c048de6..991a05828d0 100644
--- a/x/jwt_helper.go
+++ b/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)
+}