From 6c8d67ddddd18ea835058bc5b4b55c4840f83247 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Wed, 13 Jan 2021 18:07:42 +0530
Subject: [PATCH 01/31] starting work on audit logs

---
 dgraph/cmd/alpha/http.go      |  28 ++++++++
 dgraph/cmd/alpha/run.go       |  75 +++++++++++++--------
 dgraph/cmd/zero/run.go        |  32 +++++++--
 ee/audit/audit.go             | 122 ++++++++++++++++++++++++++++++++++
 ee/utils_ee.go                |   3 +
 go.mod                        |   3 +
 go.sum                        |  20 ++++++
 graphql/admin/current_user.go |  29 +-------
 worker/config.go              |   3 +
 x/config.go                   |   5 ++
 x/jwt_helper.go               |  32 +++++++++
 x/logger.go                   |  41 ++++++++++++
 12 files changed, 329 insertions(+), 64 deletions(-)
 create mode 100644 ee/audit/audit.go
 create mode 100644 x/jwt_helper.go
 create mode 100644 x/logger.go

diff --git a/dgraph/cmd/alpha/http.go b/dgraph/cmd/alpha/http.go
index 5978479bfcd..c9623dd1aef 100644
--- a/dgraph/cmd/alpha/http.go
+++ b/dgraph/cmd/alpha/http.go
@@ -31,6 +31,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/dgraph-io/dgraph/ee/audit"
+
 	"github.com/dgraph-io/dgo/v200"
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	"github.com/dgraph-io/dgraph/edgraph"
@@ -51,6 +53,32 @@ func allowed(method string) bool {
 	return method == http.MethodPost || method == http.MethodPut
 }
 
+func auditRequestWithHandler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		startTime := time.Now().UnixNano()
+		lrw := audit.NewResponseWriter(w)
+		var buf bytes.Buffer
+		tee := io.TeeReader(r.Body, &buf)
+		r.Body = ioutil.NopCloser(tee)
+		next.ServeHTTP(lrw, r)
+		r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
+		audit.Auditor.MaybeAuditFromCtx(lrw, r, startTime)
+	})
+}
+
+func auditRequest(next http.HandlerFunc) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		startTime := time.Now().UnixNano()
+		lrw := audit.NewResponseWriter(w)
+		var buf bytes.Buffer
+		tee := io.TeeReader(r.Body, &buf)
+		r.Body = ioutil.NopCloser(tee)
+		next.ServeHTTP(lrw, r)
+		r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
+		audit.Auditor.MaybeAuditFromCtx(lrw, r, startTime)
+	})
+}
+
 // Common functionality for these request handlers. Returns true if the request is completely
 // handled here and nothing further needs to be done.
 func commonHandler(w http.ResponseWriter, r *http.Request) bool {
diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 2b7c8209979..56d8ab7ea8d 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -36,6 +36,8 @@ import (
 	"time"
 
 	badgerpb "github.com/dgraph-io/badger/v3/pb"
+	"github.com/dgraph-io/dgraph/ee/audit"
+
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	"github.com/dgraph-io/dgraph/edgraph"
 	"github.com/dgraph-io/dgraph/ee/enc"
@@ -196,6 +198,10 @@ they form a Raft group and provide synchronous replication.
 		`Cache percentages summing up to 100 for various caches (FORMAT:
 		PostingListCache,PstoreBlockCache,PstoreIndexCache,WAL).`)
 
+	flag.Bool("audit_enabled", false, "Set to true to enable audit logs.")
+	// todo(aman): check what to set audit_dir default to
+	flag.String("audit_dir", "./", "Set path to directory where to save the audit logs.")
+
 	// TLS configurations
 	x.RegisterServerTLSFlags(flag)
 }
@@ -422,18 +428,19 @@ func setupServer(closer *z.Closer) {
 		log.Fatal(err)
 	}
 
-	http.HandleFunc("/query", queryHandler)
-	http.HandleFunc("/query/", queryHandler)
-	http.HandleFunc("/mutate", mutationHandler)
-	http.HandleFunc("/mutate/", mutationHandler)
-	http.HandleFunc("/commit", commitHandler)
-	http.HandleFunc("/alter", alterHandler)
-	http.HandleFunc("/health", healthCheck)
-	http.HandleFunc("/state", stateHandler)
-	http.HandleFunc("/jemalloc", x.JemallocHandler)
+	// http.HandleWithMidd("", mainhalder, middlewares....)
+	http.Handle("/query", auditRequest(queryHandler))
+	http.Handle("/query/", auditRequest(queryHandler))
+	http.Handle("/mutate", auditRequest(mutationHandler))
+	http.Handle("/mutate/", auditRequest(mutationHandler))
+	http.Handle("/commit", auditRequest(commitHandler))
+	http.Handle("/alter", auditRequest(alterHandler))
+	http.Handle("/health", auditRequest(healthCheck))
+	http.Handle("/state", auditRequest(stateHandler))
+	http.Handle("/jemalloc", auditRequest(x.JemallocHandler))
 
 	// TODO: Figure out what this is for?
-	http.HandleFunc("/debug/store", storeStatsHandler)
+	http.Handle("/debug/store", auditRequest(storeStatsHandler))
 
 	introspection := Alpha.Conf.GetBool("graphql_introspection")
 
@@ -456,8 +463,8 @@ func setupServer(closer *z.Closer) {
 	var gqlHealthStore *admin.GraphQLHealthStore
 	// Do not use := notation here because adminServer is a global variable.
 	mainServer, adminServer, gqlHealthStore = admin.NewServers(introspection, &globalEpoch, closer)
-	http.Handle("/graphql", mainServer.HTTPHandler())
-	http.HandleFunc("/probe/graphql", func(w http.ResponseWriter, r *http.Request) {
+	http.Handle("/graphql", auditRequestWithHandler(mainServer.HTTPHandler()))
+	http.Handle("/probe/graphql", auditRequest(func(w http.ResponseWriter, r *http.Request) {
 		healthStatus := gqlHealthStore.GetHealth()
 		httpStatusCode := http.StatusOK
 		if !healthStatus.Healthy {
@@ -467,19 +474,20 @@ func setupServer(closer *z.Closer) {
 		w.Header().Set("Content-Type", "application/json")
 		x.Check2(w.Write([]byte(fmt.Sprintf(`{"status":"%s","schemaUpdateCounter":%d}`,
 			healthStatus.StatusMsg, atomic.LoadUint64(&globalEpoch)))))
-	})
-	http.Handle("/admin", allowedMethodsHandler(allowedMethods{
+	}))
+	http.Handle("/admin", auditRequestWithHandler(allowedMethodsHandler(allowedMethods{
 		http.MethodGet:     true,
 		http.MethodPost:    true,
 		http.MethodOptions: true,
-	}, adminAuthHandler(adminServer.HTTPHandler())))
+	}, adminAuthHandler(adminServer.HTTPHandler()))))
 
-	http.Handle("/admin/schema", adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter,
+	http.Handle("/admin/schema", auditRequestWithHandler(adminAuthHandler(http.HandlerFunc(func(
+		w http.ResponseWriter,
 		r *http.Request) {
 		adminSchemaHandler(w, r, adminServer)
-	})))
+	}))))
 
-	http.Handle("/admin/schema/validate", http.HandlerFunc(func(w http.ResponseWriter,
+	http.Handle("/admin/schema/validate", auditRequest(http.HandlerFunc(func(w http.ResponseWriter,
 		r *http.Request) {
 		schema := readRequest(w, r)
 		w.Header().Set("Content-Type", "application/json")
@@ -494,31 +502,33 @@ func setupServer(closer *z.Closer) {
 		w.WriteHeader(http.StatusBadRequest)
 		errs := strings.Split(strings.TrimSpace(err.Error()), "\n")
 		x.SetStatusWithErrors(w, x.ErrorInvalidRequest, errs)
-	}))
+	})))
 
-	http.Handle("/admin/shutdown", allowedMethodsHandler(allowedMethods{http.MethodGet: true},
+	http.Handle("/admin/shutdown", auditRequestWithHandler(allowedMethodsHandler(allowedMethods{http.
+		MethodGet: true},
 		adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			shutDownHandler(w, r, adminServer)
-		}))))
+		})))))
 
-	http.Handle("/admin/draining", allowedMethodsHandler(allowedMethods{
+	http.Handle("/admin/draining", auditRequestWithHandler(allowedMethodsHandler(allowedMethods{
 		http.MethodPut:  true,
 		http.MethodPost: true,
 	}, adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		drainingHandler(w, r, adminServer)
-	}))))
+	})))))
 
-	http.Handle("/admin/export", allowedMethodsHandler(allowedMethods{http.MethodGet: true},
+	http.Handle("/admin/export", auditRequestWithHandler(allowedMethodsHandler(
+		allowedMethods{http.MethodGet: true},
 		adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			exportHandler(w, r, adminServer)
-		}))))
+		})))))
 
-	http.Handle("/admin/config/cache_mb", allowedMethodsHandler(allowedMethods{
+	http.Handle("/admin/config/cache_mb", auditRequestWithHandler(allowedMethodsHandler(allowedMethods{
 		http.MethodGet: true,
 		http.MethodPut: true,
 	}, adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		memoryLimitHandler(w, r, adminServer)
-	}))))
+	})))))
 
 	addr := fmt.Sprintf("%s:%d", laddr, httpPort())
 	glog.Infof("Bringing up GraphQL HTTP API at %s/graphql", addr)
@@ -527,8 +537,8 @@ func setupServer(closer *z.Closer) {
 	// Add OpenCensus z-pages.
 	zpages.Handle(http.DefaultServeMux, "/z")
 
-	http.HandleFunc("/", homeHandler)
-	http.HandleFunc("/ui/keywords", keywordHandler)
+	http.Handle("/", auditRequest(homeHandler))
+	http.Handle("/ui/keywords", auditRequest(keywordHandler))
 
 	// Initialize the servers.
 	admin.ServerCloser.AddRunning(3)
@@ -602,6 +612,8 @@ func run() {
 
 		MutationsMode: worker.AllowMutations,
 		AuthToken:     Alpha.Conf.GetString("auth_token"),
+		AuditEnabled:  Alpha.Conf.GetBool("audit_enabled"),
+		AuditDir:      Alpha.Conf.GetString("audit_dir"),
 	}
 
 	secretFile := Alpha.Conf.GetString("acl_secret_file")
@@ -663,6 +675,8 @@ func run() {
 		LudicrousConcurrency: Alpha.Conf.GetInt("ludicrous_concurrency"),
 		TLSClientConfig:      tlsClientConf,
 		TLSServerConfig:      tlsServerConf,
+		AuditEnabled:         Alpha.Conf.GetBool("audit_enabled"),
+		HmacSecret:           opts.HmacSecret,
 	}
 	x.WorkerConfig.Parse(Alpha.Conf)
 	x.CheckFlag(x.WorkerConfig.Raft, "group", "idx", "learner")
@@ -785,6 +799,7 @@ func run() {
 	// close alpha. This closer is for closing and waiting that subscription.
 	adminCloser := z.NewCloser(1)
 
+	x.Check(audit.InitAuditor())
 	setupServer(adminCloser)
 	glog.Infoln("GRPC and HTTP stopped.")
 
@@ -798,6 +813,8 @@ func run() {
 	adminCloser.SignalAndWait()
 	glog.Infoln("adminCloser closed.")
 
+	audit.Auditor.Close()
+
 	worker.State.Dispose()
 	x.RemoveCidFile()
 	glog.Info("worker.State disposed.")
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index 66b2ad0d6bf..8351be8dfd9 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -17,9 +17,12 @@
 package zero
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"log"
 	"net"
 	"net/http"
@@ -28,6 +31,8 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/dgraph-io/dgraph/ee/audit"
+
 	"go.opencensus.io/plugin/ocgrpc"
 	otrace "go.opencensus.io/trace"
 	"go.opencensus.io/zpages"
@@ -254,13 +259,26 @@ func run() {
 	x.Check(err)
 	go x.StartListenHttpAndHttps(httpListener, tlsCfg, st.zero.closer)
 
-	http.HandleFunc("/health", st.pingResponse)
-	http.HandleFunc("/state", st.getState)
-	http.HandleFunc("/removeNode", st.removeNode)
-	http.HandleFunc("/moveTablet", st.moveTablet)
-	http.HandleFunc("/assign", st.assign)
-	http.HandleFunc("/enterpriseLicense", st.applyEnterpriseLicense)
-	http.HandleFunc("/jemalloc", x.JemallocHandler)
+	auditRequest := func(next http.HandlerFunc) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			startTime := time.Now().UnixNano()
+			lrw := audit.NewResponseWriter(w)
+			var buf bytes.Buffer
+			tee := io.TeeReader(r.Body, &buf)
+			r.Body = ioutil.NopCloser(tee)
+			next.ServeHTTP(lrw, r)
+			r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
+			audit.Auditor.MaybeAuditFromCtx(lrw, r, startTime)
+		})
+	}
+
+	http.Handle("/health", auditRequest(st.pingResponse))
+	http.Handle("/state", auditRequest(st.getState))
+	http.Handle("/removeNode", auditRequest(st.removeNode))
+	http.Handle("/moveTablet", auditRequest(st.moveTablet))
+	http.Handle("/assign", auditRequest(st.assign))
+	http.Handle("/enterpriseLicense", auditRequest(st.applyEnterpriseLicense))
+	http.Handle("/jemalloc", auditRequest(x.JemallocHandler))
 	zpages.Handle(http.DefaultServeMux, "/z")
 
 	// This must be here. It does not work if placed before Grpc init.
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
new file mode 100644
index 00000000000..d5111440f77
--- /dev/null
+++ b/ee/audit/audit.go
@@ -0,0 +1,122 @@
+package audit
+
+import (
+	"encoding/json"
+	"errors"
+	"io/ioutil"
+	"log"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/golang/glog"
+
+	"github.com/dgraph-io/dgraph/worker"
+	"github.com/dgraph-io/dgraph/x"
+)
+
+type AuditEvent struct {
+	User        string              `json:"user"`
+	ServerHost  string              `json:"server_host"`
+	ClientHost  string              `json:"client_host"`
+	Timestamp   int64               `json:"timestamp"`
+	Req         string              `json:"req"`
+	Status      int                 `json:"status"`
+	QueryParams map[string][]string `json:"query_params"`
+	TimeTaken   int64               `json:"time_taken"`
+}
+
+const (
+	UnauthorisedUser = "UnauthorisedUser"
+	NoUser           = "NoUser"
+	UnknownUser      = "UnknownUser"
+)
+
+var Auditor *AuditLogger
+
+// todo add file rotation etc
+type AuditLogger struct {
+	log            *log.Logger
+	AuditorEnabled bool
+}
+
+func InitAuditor() error {
+	if !worker.EnterpriseEnabled() {
+		return errors.New("enterprise features are disabled. You can enable them by " +
+			"supplying the appropriate license file to Dgraph Zero using the HTTP endpoint")
+	}
+
+	e, err := os.OpenFile(worker.Config.AuditDir+"/dgraph_audit.json",
+		os.O_WRONLY|os.O_CREATE|os.O_APPEND,
+		0666)
+	if err != nil {
+		return err
+	}
+
+	Auditor = &AuditLogger{
+		log:            log.New(e, "", 0),
+		AuditorEnabled: true,
+	}
+
+	//Auditor.startListeningToLogAudit()
+	return nil
+}
+
+func (a *AuditLogger) Close() {
+	glog.Infof("Closing auditor")
+}
+
+func (a *AuditLogger) MaybeAuditFromCtx(w *ResponseWriter, r *http.Request, startTime int64) {
+	var userId string
+	var err error
+	token := r.Header.Get("X-Dgraph-AccessToken")
+	if token == "" {
+		if x.WorkerConfig.AclEnabled {
+			userId = UnauthorisedUser
+		} else {
+			userId = NoUser
+		}
+	} else {
+		userId, err = x.ExtractUserName(token)
+		if err != nil {
+			userId = UnknownUser
+		}
+	}
+
+	all, _ := ioutil.ReadAll(r.Body)
+	event := &AuditEvent{
+		User:        userId,
+		ServerHost:  x.WorkerConfig.MyAddr,
+		ClientHost:  r.RemoteAddr,
+		Timestamp:   time.Now().Unix(),
+		Req:         string(all),
+		Status:      w.statusCode,
+		QueryParams: r.URL.Query(),
+		TimeTaken:   time.Now().UnixNano() - startTime,
+	}
+	a.MaybeAuditFromEvent(event)
+}
+
+func (a *AuditLogger) MaybeAuditFromEvent(event *AuditEvent) {
+	if !a.AuditorEnabled {
+		return
+	}
+	b, _ := json.Marshal(event)
+	a.log.Print(string(b))
+}
+
+type ResponseWriter struct {
+	http.ResponseWriter
+	statusCode int
+}
+
+func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
+	// WriteHeader(int) is not called if our response implicitly returns 200 OK, so
+	// we default to that status code.
+	return &ResponseWriter{w, http.StatusOK}
+}
+
+func (lrw *ResponseWriter) WriteHeader(code int) {
+	lrw.statusCode = code
+	lrw.ResponseWriter.WriteHeader(code)
+}
diff --git a/ee/utils_ee.go b/ee/utils_ee.go
index ad546478de7..cf2007a5c6b 100644
--- a/ee/utils_ee.go
+++ b/ee/utils_ee.go
@@ -37,5 +37,8 @@ func GetEEFeaturesList() []string {
 	} else {
 		ee = append(ee, "backup_restore")
 	}
+	if worker.Config.AuditEnabled {
+		ee = append(ee, "audit")
+	}
 	return ee
 }
diff --git a/go.mod b/go.mod
index 3ab01e09e51..e46498fafeb 100644
--- a/go.mod
+++ b/go.mod
@@ -60,6 +60,8 @@ require (
 	github.com/twpayne/go-geom v1.0.5
 	go.etcd.io/etcd v0.0.0-20190228193606-a943ad0ee4c9
 	go.opencensus.io v0.22.5
+	go.uber.org/multierr v1.6.0 // indirect
+	go.uber.org/zap v1.16.0
 	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9
 	golang.org/x/net v0.0.0-20201021035429-f5854403a974
 	golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9
@@ -68,6 +70,7 @@ require (
 	golang.org/x/tools v0.0.0-20210106214847-113979e3529a
 	google.golang.org/grpc v1.23.0
 	gopkg.in/DataDog/dd-trace-go.v1 v1.13.1 // indirect
+	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 	gopkg.in/square/go-jose.v2 v2.3.1
 	gopkg.in/yaml.v2 v2.2.4
 )
diff --git a/go.sum b/go.sum
index ea46511b781..0e1c63af36f 100644
--- a/go.sum
+++ b/go.sum
@@ -97,8 +97,10 @@ github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8Nz
 github.com/coreos/go-semver v0.2.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/coreos/go-systemd v0.0.0-20180511133405-39ca1b05acc7/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
+github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e h1:Wf6HqHfScWJN9/ZjdUKyjop4mf3Qdd+1TvvltAvM3m8=
 github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/coreos/pkg v0.0.0-20160727233714-3ac0863d7acf/go.mod h1:E3G3o1h8I7cfcXa63jLwjI0eiQQMgzzUDFVpN/nH/eA=
+github.com/coreos/pkg v0.0.0-20180928190104-399ea9e2e55f h1:lBNOc5arjvs8E5mO2tbpBpLoyyu8B6e44T7hJy6potg=
 github.com/coreos/pkg v0.0.0-20180928190104-399ea9e2e55f/go.mod h1:E3G3o1h8I7cfcXa63jLwjI0eiQQMgzzUDFVpN/nH/eA=
 github.com/cosiner/argv v0.1.0/go.mod h1:EusR6TucWKX+zFgtdUsKT2Cvg45K5rtpCcWz4hK06d8=
 github.com/couchbase/ghistogram v0.1.0/go.mod h1:s1Jhy76zqfEecpNWJfWUiKZookAFaiGOEoyzgHt9i7k=
@@ -584,10 +586,24 @@ go.opencensus.io v0.22.5 h1:dntmOdLpSpHlVqbW5Eay97DelsZHe+55D+xC6i0dDS0=
 go.opencensus.io v0.22.5/go.mod h1:5pWMHQbX5EPX2/62yrJeAkowc+lfs/XD7Uxpq3pI6kk=
 go.starlark.net v0.0.0-20190702223751-32f345186213/go.mod h1:c1/X6cHgvdXj6pUlmWKMkuqRnW4K8x2vwt6JAaaircg=
 go.uber.org/atomic v1.3.2/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
+go.uber.org/atomic v1.4.0 h1:cxzIVoETapQEqDhQu3QfnvXAV4AlzcvUCxkVUFw3+EU=
 go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
+go.uber.org/atomic v1.6.0 h1:Ezj3JGmsOnG1MoRWQkPBsKLe9DwWD9QeXzTRzzldNVk=
+go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
+go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
+go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/multierr v1.1.0 h1:HoEmRHQPVSqub6w2z2d2EOVs2fjyFRGyofhKuyDq0QI=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
+go.uber.org/multierr v1.5.0 h1:KCa4XfM8CWFCpxXRGok+Q0SS/0XBhMDbHHGABQLvD2A=
+go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU=
+go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
+go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
+go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA=
 go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
+go.uber.org/zap v1.10.0 h1:ORx85nbTijNz8ljznvCMR1ZBIPKFn3jQrag10X2AsuM=
 go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
+go.uber.org/zap v1.16.0 h1:uFRZXykJGK9lLY4HtgSw44DnIcAM+kRBP7x5m+NpAOM=
+go.uber.org/zap v1.16.0/go.mod h1:MA8QOfq0BHJwdXa996Y4dYkAqRKB8/1K1QMMZVaNZjQ=
 golang.org/x/arch v0.0.0-20190927153633-4e8777c89be4/go.mod h1:flIaEI6LNU6xOCD5PaJvn9wGP0agmIOqjrtsKGRguv4=
 golang.org/x/arch v0.0.0-20201008161808-52c3e6f60cff/go.mod h1:flIaEI6LNU6xOCD5PaJvn9wGP0agmIOqjrtsKGRguv4=
 golang.org/x/crypto v0.0.0-20180608092829-8ac0e0d97ce4/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
@@ -723,6 +739,8 @@ golang.org/x/tools v0.0.0-20190628153133-6cdbf07be9d0/go.mod h1:/rFqwRUd4F7ZHNgw
 golang.org/x/tools v0.0.0-20190816200558-6889da9d5479/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20190911174233-4f2ddba30aff/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20191012152004-8de300cfc20a/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.0.0-20191029041327-9cc4af7d6b2c/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.0.0-20191029190741-b9c20aec41a5/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20191112195655-aa38f8e97acc/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20191127201027-ecd32218bd7f/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
@@ -789,6 +807,8 @@ gopkg.in/ini.v1 v1.42.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k=
 gopkg.in/ini.v1 v1.51.0 h1:AQvPpx3LzTDM0AjnIRlVFwFFGC+npRopjZxLJj6gdno=
 gopkg.in/ini.v1 v1.51.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k=
 gopkg.in/mgo.v2 v2.0.0-20180705113604-9856a29383ce/go.mod h1:yeKp02qBN3iKW1OzL3MGk2IdtZzaj7SFntXj72NppTA=
+gopkg.in/natefinch/lumberjack.v2 v2.0.0 h1:1Lc07Kr7qY4U2YPouBjpCLxpiyxIVoxqXgkXLknAOE8=
+gopkg.in/natefinch/lumberjack.v2 v2.0.0/go.mod h1:l0ndWWf7gzL7RNwBG7wST/UCcT4T24xpD6X8LsfU/+k=
 gopkg.in/resty.v1 v1.12.0/go.mod h1:mDo4pnntr5jdWRML875a/NmxYqAlA73dVijT2AXvQQo=
 gopkg.in/square/go-jose.v2 v2.3.1 h1:SK5KegNXmKmqE342YYN2qPHEnUYeoMiXXl1poUlI+o4=
 gopkg.in/square/go-jose.v2 v2.3.1/go.mod h1:M9dMgbHiYLoDGQrXy7OpJDJWiKiU//h+vD76mk0e1AI=
diff --git a/graphql/admin/current_user.go b/graphql/admin/current_user.go
index d7e92b1c854..6bbcc43ee13 100644
--- a/graphql/admin/current_user.go
+++ b/graphql/admin/current_user.go
@@ -22,10 +22,7 @@ import (
 	"github.com/dgraph-io/dgraph/gql"
 	"github.com/dgraph-io/dgraph/graphql/resolve"
 	"github.com/dgraph-io/dgraph/graphql/schema"
-	"github.com/dgraph-io/dgraph/worker"
 	"github.com/dgraph-io/dgraph/x"
-	"github.com/dgrijalva/jwt-go"
-	"github.com/pkg/errors"
 )
 
 type currentUserResolver struct {
@@ -38,31 +35,7 @@ func extractName(ctx context.Context) (string, error) {
 		return "", err
 	}
 
-	// Code copied from access_ee.go. Couldn't put the code in x, because of dependency on
-	// worker. (worker.Config.HmacSecret)
-	token, err := jwt.Parse(accessJwt[0], func(token *jwt.Token) (interface{}, error) {
-		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
-			return nil, errors.Errorf("unexpected signing method: %v",
-				token.Header["alg"])
-		}
-		return []byte(worker.Config.HmacSecret), nil
-	})
-
-	if err != nil {
-		return "", errors.Wrapf(err, "unable to parse jwt token")
-	}
-
-	claims, ok := token.Claims.(jwt.MapClaims)
-	if !ok || !token.Valid {
-		return "", errors.Errorf("claims in jwt token is not map claims")
-	}
-
-	userId, ok := claims["userid"].(string)
-	if !ok {
-		return "", errors.Errorf("userid in claims is not a string:%v", userId)
-	}
-
-	return userId, nil
+	return x.ExtractUserName(accessJwt[0])
 }
 
 func (gsr *currentUserResolver) Rewrite(ctx context.Context,
diff --git a/worker/config.go b/worker/config.go
index ada465cee26..e0b6076395a 100644
--- a/worker/config.go
+++ b/worker/config.go
@@ -70,6 +70,9 @@ type Options struct {
 	CachePercentage string
 	// CacheMb is the total memory allocated between all the caches.
 	CacheMb int64
+
+	AuditDir     string
+	AuditEnabled bool
 }
 
 // Config holds an instance of the server options..
diff --git a/x/config.go b/x/config.go
index 17bfe5bbcdf..b030056e7f5 100644
--- a/x/config.go
+++ b/x/config.go
@@ -87,6 +87,8 @@ type WorkerOptions struct {
 	StrictMutations bool
 	// AclEnabled indicates whether the enterprise ACL feature is turned on.
 	AclEnabled bool
+	// HmacSecret stores the secret used to sign JSON Web Tokens (JWT).
+	HmacSecret SensitiveByteSlice
 	// AbortOlderThan tells Dgraph to discard transactions that are older than this duration.
 	AbortOlderThan time.Duration
 	// SnapshotAfter indicates the number of entries in the RAFT logs that are needed
@@ -110,6 +112,9 @@ type WorkerOptions struct {
 	LogRequest int32
 	// If true, we should call msync or fsync after every write to survive hard reboots.
 	HardSync bool
+
+	// AuditEnabled determines whether the audit feature is enabled or not
+	AuditEnabled bool
 }
 
 // WorkerConfig stores the global instance of the worker package's options.
diff --git a/x/jwt_helper.go b/x/jwt_helper.go
new file mode 100644
index 00000000000..34efa2f9b9e
--- /dev/null
+++ b/x/jwt_helper.go
@@ -0,0 +1,32 @@
+package x
+
+import (
+	"github.com/dgrijalva/jwt-go"
+	"github.com/pkg/errors"
+)
+
+func ExtractUserName(jwtToken string) (string, error) {
+	token, err := jwt.Parse(jwtToken, func(token *jwt.Token) (interface{}, error) {
+		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
+			return nil, errors.Errorf("unexpected signing method: %v",
+				token.Header["alg"])
+		}
+		return []byte(WorkerConfig.HmacSecret), nil
+	})
+
+	if err != nil {
+		return "", errors.Wrapf(err, "unable to parse jwt token")
+	}
+
+	claims, ok := token.Claims.(jwt.MapClaims)
+	if !ok || !token.Valid {
+		return "", errors.Errorf("claims in jwt token is not map claims")
+	}
+
+	userId, ok := claims["userid"].(string)
+	if !ok {
+		return "", errors.Errorf("userid in claims is not a string:%v", userId)
+	}
+
+	return userId, nil
+}
diff --git a/x/logger.go b/x/logger.go
new file mode 100644
index 00000000000..2ead08f21e9
--- /dev/null
+++ b/x/logger.go
@@ -0,0 +1,41 @@
+package x
+
+import (
+	"os"
+
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+	"golang.org/x/net/context"
+)
+
+var logging LoggerImpl
+
+type ILogger interface {
+	Audit(ctx context.Context, args ...interface{})
+	Sync()
+}
+
+func initLogger() error {
+	core := zapcore.NewCore(zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
+		zapcore.AddSync(os.Stdout), zap.DebugLevel)
+
+	logger := zap.New(core)
+	sugarLogger := logger.Sugar()
+
+	logging = LoggerImpl{
+		sugarLogger: sugarLogger,
+	}
+	return nil
+}
+
+type LoggerImpl struct {
+	sugarLogger *zap.SugaredLogger
+}
+
+func Audit(args ...interface{}) {
+	logging.sugarLogger.Info(args...)
+}
+
+func (l *LoggerImpl) Sync() {
+	_ = l.sugarLogger.Sync()
+}

From c77f3d03dd02356b88851d2070139ce6235ddcd0 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Thu, 14 Jan 2021 15:13:36 +0530
Subject: [PATCH 02/31] making audit logs to work with alpha

---
 dgraph/cmd/alpha/http.go  |  28 ---------
 dgraph/cmd/alpha/run.go   |  55 +++++++++--------
 dgraph/cmd/live/batch.go  |   3 +-
 dgraph/cmd/zero/oracle.go |   5 +-
 dgraph/cmd/zero/run.go    |  37 ++++-------
 edgraph/server.go         |   5 +-
 ee/audit/audit.go         | 125 +++++++++++++++++++-------------------
 ee/audit/interceptor.go   | 103 +++++++++++++++++++++++++++++++
 posting/list.go           |   3 +-
 worker/draft.go           |   3 +-
 x/logger.go               |  52 ++++++++++------
 x/x.go                    |   2 +
 12 files changed, 248 insertions(+), 173 deletions(-)
 create mode 100644 ee/audit/interceptor.go

diff --git a/dgraph/cmd/alpha/http.go b/dgraph/cmd/alpha/http.go
index c9623dd1aef..5978479bfcd 100644
--- a/dgraph/cmd/alpha/http.go
+++ b/dgraph/cmd/alpha/http.go
@@ -31,8 +31,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/dgraph-io/dgraph/ee/audit"
-
 	"github.com/dgraph-io/dgo/v200"
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	"github.com/dgraph-io/dgraph/edgraph"
@@ -53,32 +51,6 @@ func allowed(method string) bool {
 	return method == http.MethodPost || method == http.MethodPut
 }
 
-func auditRequestWithHandler(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		startTime := time.Now().UnixNano()
-		lrw := audit.NewResponseWriter(w)
-		var buf bytes.Buffer
-		tee := io.TeeReader(r.Body, &buf)
-		r.Body = ioutil.NopCloser(tee)
-		next.ServeHTTP(lrw, r)
-		r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
-		audit.Auditor.MaybeAuditFromCtx(lrw, r, startTime)
-	})
-}
-
-func auditRequest(next http.HandlerFunc) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		startTime := time.Now().UnixNano()
-		lrw := audit.NewResponseWriter(w)
-		var buf bytes.Buffer
-		tee := io.TeeReader(r.Body, &buf)
-		r.Body = ioutil.NopCloser(tee)
-		next.ServeHTTP(lrw, r)
-		r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
-		audit.Auditor.MaybeAuditFromCtx(lrw, r, startTime)
-	})
-}
-
 // Common functionality for these request handlers. Returns true if the request is completely
 // handled here and nothing further needs to be done.
 func commonHandler(w http.ResponseWriter, r *http.Request) bool {
diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 56d8ab7ea8d..9ee15a4055c 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -390,6 +390,7 @@ func serveGRPC(l net.Listener, tlsCfg *tls.Config, closer *z.Closer) {
 		grpc.MaxSendMsgSize(x.GrpcMaxSize),
 		grpc.MaxConcurrentStreams(1000),
 		grpc.StatsHandler(&ocgrpc.ServerHandler{}),
+		grpc.UnaryInterceptor(audit.AuditRequestGRPC),
 	}
 	if tlsCfg != nil {
 		opt = append(opt, grpc.Creds(credentials.NewTLS(tlsCfg)))
@@ -428,19 +429,18 @@ func setupServer(closer *z.Closer) {
 		log.Fatal(err)
 	}
 
-	// http.HandleWithMidd("", mainhalder, middlewares....)
-	http.Handle("/query", auditRequest(queryHandler))
-	http.Handle("/query/", auditRequest(queryHandler))
-	http.Handle("/mutate", auditRequest(mutationHandler))
-	http.Handle("/mutate/", auditRequest(mutationHandler))
-	http.Handle("/commit", auditRequest(commitHandler))
-	http.Handle("/alter", auditRequest(alterHandler))
-	http.Handle("/health", auditRequest(healthCheck))
-	http.Handle("/state", auditRequest(stateHandler))
-	http.Handle("/jemalloc", auditRequest(x.JemallocHandler))
+	http.Handle("/query", audit.AuditRequestHttp(http.HandlerFunc(queryHandler)))
+	http.Handle("/query/", audit.AuditRequestHttp(http.HandlerFunc(queryHandler)))
+	http.Handle("/mutate", audit.AuditRequestHttp(http.HandlerFunc(mutationHandler)))
+	http.Handle("/mutate/", audit.AuditRequestHttp(http.HandlerFunc(mutationHandler)))
+	http.Handle("/commit", audit.AuditRequestHttp(http.HandlerFunc(commitHandler)))
+	http.Handle("/alter", audit.AuditRequestHttp(http.HandlerFunc(alterHandler)))
+	http.HandleFunc("/health", healthCheck)
+	http.HandleFunc("/state", stateHandler)
+	http.HandleFunc("/jemalloc", x.JemallocHandler)
 
 	// TODO: Figure out what this is for?
-	http.Handle("/debug/store", auditRequest(storeStatsHandler))
+	http.HandleFunc("/debug/store", storeStatsHandler)
 
 	introspection := Alpha.Conf.GetBool("graphql_introspection")
 
@@ -463,8 +463,9 @@ func setupServer(closer *z.Closer) {
 	var gqlHealthStore *admin.GraphQLHealthStore
 	// Do not use := notation here because adminServer is a global variable.
 	mainServer, adminServer, gqlHealthStore = admin.NewServers(introspection, &globalEpoch, closer)
-	http.Handle("/graphql", auditRequestWithHandler(mainServer.HTTPHandler()))
-	http.Handle("/probe/graphql", auditRequest(func(w http.ResponseWriter, r *http.Request) {
+	http.Handle("/graphql", audit.AuditRequestHttp(mainServer.HTTPHandler()))
+	http.Handle("/probe/graphql", audit.AuditRequestHttp(http.HandlerFunc(func(w http.ResponseWriter,
+		r *http.Request) {
 		healthStatus := gqlHealthStore.GetHealth()
 		httpStatusCode := http.StatusOK
 		if !healthStatus.Healthy {
@@ -474,20 +475,20 @@ func setupServer(closer *z.Closer) {
 		w.Header().Set("Content-Type", "application/json")
 		x.Check2(w.Write([]byte(fmt.Sprintf(`{"status":"%s","schemaUpdateCounter":%d}`,
 			healthStatus.StatusMsg, atomic.LoadUint64(&globalEpoch)))))
-	}))
-	http.Handle("/admin", auditRequestWithHandler(allowedMethodsHandler(allowedMethods{
+	})))
+	http.Handle("/admin", audit.AuditRequestHttp(allowedMethodsHandler(allowedMethods{
 		http.MethodGet:     true,
 		http.MethodPost:    true,
 		http.MethodOptions: true,
 	}, adminAuthHandler(adminServer.HTTPHandler()))))
 
-	http.Handle("/admin/schema", auditRequestWithHandler(adminAuthHandler(http.HandlerFunc(func(
+	http.Handle("/admin/schema", audit.AuditRequestHttp(adminAuthHandler(http.HandlerFunc(func(
 		w http.ResponseWriter,
 		r *http.Request) {
 		adminSchemaHandler(w, r, adminServer)
 	}))))
 
-	http.Handle("/admin/schema/validate", auditRequest(http.HandlerFunc(func(w http.ResponseWriter,
+	http.Handle("/admin/schema/validate", audit.AuditRequestHttp(http.HandlerFunc(func(w http.ResponseWriter,
 		r *http.Request) {
 		schema := readRequest(w, r)
 		w.Header().Set("Content-Type", "application/json")
@@ -504,26 +505,26 @@ func setupServer(closer *z.Closer) {
 		x.SetStatusWithErrors(w, x.ErrorInvalidRequest, errs)
 	})))
 
-	http.Handle("/admin/shutdown", auditRequestWithHandler(allowedMethodsHandler(allowedMethods{http.
+	http.Handle("/admin/shutdown", audit.AuditRequestHttp(allowedMethodsHandler(allowedMethods{http.
 		MethodGet: true},
 		adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			shutDownHandler(w, r, adminServer)
 		})))))
 
-	http.Handle("/admin/draining", auditRequestWithHandler(allowedMethodsHandler(allowedMethods{
+	http.Handle("/admin/draining", audit.AuditRequestHttp(allowedMethodsHandler(allowedMethods{
 		http.MethodPut:  true,
 		http.MethodPost: true,
 	}, adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		drainingHandler(w, r, adminServer)
 	})))))
 
-	http.Handle("/admin/export", auditRequestWithHandler(allowedMethodsHandler(
+	http.Handle("/admin/export", audit.AuditRequestHttp(allowedMethodsHandler(
 		allowedMethods{http.MethodGet: true},
 		adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			exportHandler(w, r, adminServer)
 		})))))
 
-	http.Handle("/admin/config/cache_mb", auditRequestWithHandler(allowedMethodsHandler(allowedMethods{
+	http.Handle("/admin/config/cache_mb", audit.AuditRequestHttp(allowedMethodsHandler(allowedMethods{
 		http.MethodGet: true,
 		http.MethodPut: true,
 	}, adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -537,8 +538,8 @@ func setupServer(closer *z.Closer) {
 	// Add OpenCensus z-pages.
 	zpages.Handle(http.DefaultServeMux, "/z")
 
-	http.Handle("/", auditRequest(homeHandler))
-	http.Handle("/ui/keywords", auditRequest(keywordHandler))
+	http.Handle("/", audit.AuditRequestHttp(http.HandlerFunc(homeHandler)))
+	http.Handle("/ui/keywords", audit.AuditRequestHttp(http.HandlerFunc(keywordHandler)))
 
 	// Initialize the servers.
 	admin.ServerCloser.AddRunning(3)
@@ -799,7 +800,11 @@ func run() {
 	// close alpha. This closer is for closing and waiting that subscription.
 	adminCloser := z.NewCloser(1)
 
-	x.Check(audit.InitAuditor())
+	// Audit is enterprise feature. If enabled, audit logs will be generate in the audit directory.
+	if opts.AuditEnabled {
+		go audit.InitAuditorIfNecessary(opts.AuditDir)
+	}
+
 	setupServer(adminCloser)
 	glog.Infoln("GRPC and HTTP stopped.")
 
@@ -813,7 +818,7 @@ func run() {
 	adminCloser.SignalAndWait()
 	glog.Infoln("adminCloser closed.")
 
-	audit.Auditor.Close()
+	audit.Close()
 
 	worker.State.Dispose()
 	x.RemoveCidFile()
diff --git a/dgraph/cmd/live/batch.go b/dgraph/cmd/live/batch.go
index d26386e1e29..2a56087e536 100644
--- a/dgraph/cmd/live/batch.go
+++ b/dgraph/cmd/live/batch.go
@@ -34,7 +34,6 @@ import (
 	"github.com/dgraph-io/badger/v3"
 	"github.com/dgraph-io/dgo/v200"
 	"github.com/dgraph-io/dgo/v200/protos/api"
-	"github.com/dgraph-io/dgraph/dgraph/cmd/zero"
 	"github.com/dgraph-io/dgraph/gql"
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/tok"
@@ -132,7 +131,7 @@ func handleError(err error, isRetry bool) {
 		dur := time.Duration(1+rand.Intn(10)) * time.Minute
 		fmt.Printf("Server is overloaded. Will retry after %s.\n", dur.Round(time.Minute))
 		time.Sleep(dur)
-	case err != zero.ErrConflict && err != dgo.ErrAborted:
+	case err != x.ErrConflict && err != dgo.ErrAborted:
 		fmt.Printf("Error while mutating: %v s.Code %v\n", s.Message(), s.Code())
 	}
 }
diff --git a/dgraph/cmd/zero/oracle.go b/dgraph/cmd/zero/oracle.go
index 24771b9abf3..056815674e9 100644
--- a/dgraph/cmd/zero/oracle.go
+++ b/dgraph/cmd/zero/oracle.go
@@ -134,7 +134,7 @@ func (o *Oracle) commit(src *api.TxnContext) error {
 	defer o.Unlock()
 
 	if o.hasConflict(src) {
-		return ErrConflict
+		return x.ErrConflict
 	}
 	// We store src.Keys as string to ensure compatibility with all the various language clients we
 	// have. But, really they are just uint64s encoded as strings. We use base 36 during creation of
@@ -310,9 +310,6 @@ func (o *Oracle) MaxPending() uint64 {
 	return o.maxAssigned
 }
 
-// ErrConflict is returned when commit couldn't succeed due to conflicts.
-var ErrConflict = errors.New("Transaction conflict")
-
 // proposeTxn proposes a txn update, and then updates src to reflect the state
 // of the commit after proposal is run.
 func (s *Server) proposeTxn(ctx context.Context, src *api.TxnContext) error {
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index 8351be8dfd9..9a76e23f875 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -17,12 +17,9 @@
 package zero
 
 import (
-	"bytes"
 	"context"
 	"crypto/tls"
 	"fmt"
-	"io"
-	"io/ioutil"
 	"log"
 	"net"
 	"net/http"
@@ -31,8 +28,6 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/dgraph-io/dgraph/ee/audit"
-
 	"go.opencensus.io/plugin/ocgrpc"
 	otrace "go.opencensus.io/trace"
 	"go.opencensus.io/zpages"
@@ -100,6 +95,11 @@ instances to achieve high-availability.
 	flag.StringP("wal", "w", "zw", "Directory storing WAL.")
 	flag.Duration("rebalance_interval", 8*time.Minute, "Interval for trying a predicate move.")
 	flag.String("enterprise_license", "", "Path to the enterprise license file.")
+
+	flag.Bool("audit_enabled", false, "Set to true to enable audit logs.")
+	// todo(aman): check what to set audit_dir default to
+	flag.String("audit_dir", "./", "Set path to directory where to save the audit logs.")
+
 	// TLS configurations
 	x.RegisterServerTLSFlags(flag)
 }
@@ -259,26 +259,13 @@ func run() {
 	x.Check(err)
 	go x.StartListenHttpAndHttps(httpListener, tlsCfg, st.zero.closer)
 
-	auditRequest := func(next http.HandlerFunc) http.Handler {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			startTime := time.Now().UnixNano()
-			lrw := audit.NewResponseWriter(w)
-			var buf bytes.Buffer
-			tee := io.TeeReader(r.Body, &buf)
-			r.Body = ioutil.NopCloser(tee)
-			next.ServeHTTP(lrw, r)
-			r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
-			audit.Auditor.MaybeAuditFromCtx(lrw, r, startTime)
-		})
-	}
-
-	http.Handle("/health", auditRequest(st.pingResponse))
-	http.Handle("/state", auditRequest(st.getState))
-	http.Handle("/removeNode", auditRequest(st.removeNode))
-	http.Handle("/moveTablet", auditRequest(st.moveTablet))
-	http.Handle("/assign", auditRequest(st.assign))
-	http.Handle("/enterpriseLicense", auditRequest(st.applyEnterpriseLicense))
-	http.Handle("/jemalloc", auditRequest(x.JemallocHandler))
+	http.HandleFunc("/health", st.pingResponse)
+	http.HandleFunc("/state", st.getState)
+	http.HandleFunc("/removeNode", st.removeNode)
+	http.HandleFunc("/moveTablet", st.moveTablet)
+	http.HandleFunc("/assign", st.assign)
+	http.HandleFunc("/enterpriseLicense", st.applyEnterpriseLicense)
+	http.HandleFunc("/jemalloc", x.JemallocHandler)
 	zpages.Handle(http.DefaultServeMux, "/z")
 
 	// This must be here. It does not work if placed before Grpc init.
diff --git a/edgraph/server.go b/edgraph/server.go
index dfa69243416..579effead7b 100644
--- a/edgraph/server.go
+++ b/edgraph/server.go
@@ -46,7 +46,6 @@ import (
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	"github.com/dgraph-io/dgraph/chunker"
 	"github.com/dgraph-io/dgraph/conn"
-	"github.com/dgraph-io/dgraph/dgraph/cmd/zero"
 	"github.com/dgraph-io/dgraph/ee"
 	"github.com/dgraph-io/dgraph/gql"
 	"github.com/dgraph-io/dgraph/posting"
@@ -571,7 +570,7 @@ func (s *Server) doMutate(ctx context.Context, qc *queryContext, resp *api.Respo
 	}
 	if !qc.req.CommitNow {
 		calculateMutationMetrics()
-		if err == zero.ErrConflict {
+		if err == x.ErrConflict {
 			err = status.Error(codes.FailedPrecondition, err.Error())
 		}
 
@@ -590,7 +589,7 @@ func (s *Server) doMutate(ctx context.Context, qc *queryContext, resp *api.Respo
 		resp.Txn.Aborted = true
 		_, _ = worker.CommitOverNetwork(ctx, resp.Txn)
 
-		if err == zero.ErrConflict {
+		if err == x.ErrConflict {
 			// We have already aborted the transaction, so the error message should reflect that.
 			return dgo.ErrAborted
 		}
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index d5111440f77..3011160e26a 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -1,12 +1,9 @@
 package audit
 
 import (
-	"encoding/json"
-	"errors"
-	"io/ioutil"
-	"log"
 	"net/http"
-	"os"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/golang/glog"
@@ -15,11 +12,13 @@ import (
 	"github.com/dgraph-io/dgraph/x"
 )
 
+var auditEnabled uint32
+
 type AuditEvent struct {
 	User        string              `json:"user"`
 	ServerHost  string              `json:"server_host"`
 	ClientHost  string              `json:"client_host"`
-	Timestamp   int64               `json:"timestamp"`
+	Endpoint    string              `json:"endpoint"`
 	Req         string              `json:"req"`
 	Status      int                 `json:"status"`
 	QueryParams map[string][]string `json:"query_params"`
@@ -28,81 +27,81 @@ type AuditEvent struct {
 
 const (
 	UnauthorisedUser = "UnauthorisedUser"
-	NoUser           = "NoUser"
 	UnknownUser      = "UnknownUser"
 )
 
-var Auditor *AuditLogger
+var auditor *auditLogger
 
-// todo add file rotation etc
-type AuditLogger struct {
-	log            *log.Logger
-	AuditorEnabled bool
+type auditLogger struct {
+	mu   sync.RWMutex
+	tick *time.Ticker
+	log  x.ILogger
 }
 
-func InitAuditor() error {
-	if !worker.EnterpriseEnabled() {
-		return errors.New("enterprise features are disabled. You can enable them by " +
-			"supplying the appropriate license file to Dgraph Zero using the HTTP endpoint")
-	}
-
-	e, err := os.OpenFile(worker.Config.AuditDir+"/dgraph_audit.json",
-		os.O_WRONLY|os.O_CREATE|os.O_APPEND,
-		0666)
-	if err != nil {
-		return err
+func InitAuditorIfNecessary(dir string) {
+	auditor = &auditLogger{
+		mu: sync.RWMutex{},
+		tick: time.NewTicker(time.Minute * 5),
 	}
 
-	Auditor = &AuditLogger{
-		log:            log.New(e, "", 0),
-		AuditorEnabled: true,
+	if !worker.EnterpriseEnabled() {
+		return
 	}
 
-	//Auditor.startListeningToLogAudit()
-	return nil
-}
-
-func (a *AuditLogger) Close() {
-	glog.Infof("Closing auditor")
-}
-
-func (a *AuditLogger) MaybeAuditFromCtx(w *ResponseWriter, r *http.Request, startTime int64) {
-	var userId string
-	var err error
-	token := r.Header.Get("X-Dgraph-AccessToken")
-	if token == "" {
-		if x.WorkerConfig.AclEnabled {
-			userId = UnauthorisedUser
-		} else {
-			userId = NoUser
-		}
-	} else {
-		userId, err = x.ExtractUserName(token)
+	initlog := func() x.ILogger{
+		logger, err := x.InitLogger(dir, "dgraph_audit.log")
 		if err != nil {
-			userId = UnknownUser
+			glog.Errorf("error while initiating auditor %v", err)
+			return nil
 		}
+		return logger
 	}
-
-	all, _ := ioutil.ReadAll(r.Body)
-	event := &AuditEvent{
-		User:        userId,
-		ServerHost:  x.WorkerConfig.MyAddr,
-		ClientHost:  r.RemoteAddr,
-		Timestamp:   time.Now().Unix(),
-		Req:         string(all),
-		Status:      w.statusCode,
-		QueryParams: r.URL.Query(),
-		TimeTaken:   time.Now().UnixNano() - startTime,
+	auditor.log = initlog()
+	for {
+		select {
+		case <- auditor.tick.C:
+			if !worker.EnterpriseEnabled() {
+				if atomic.LoadUint32(&auditEnabled) != 0 {
+					atomic.StoreUint32(&auditEnabled, 0)
+					auditor.mu.Lock()
+					auditor.log = nil
+					auditor.mu.Unlock()
+					continue
+				}
+			}
+
+			if atomic.LoadUint32(&auditEnabled) != 1 {
+				atomic.StoreUint32(&auditEnabled, 1)
+				auditor.mu.Lock()
+				auditor.log = initlog()
+				auditor.mu.Unlock()
+			}
+		}
 	}
-	a.MaybeAuditFromEvent(event)
 }
 
-func (a *AuditLogger) MaybeAuditFromEvent(event *AuditEvent) {
-	if !a.AuditorEnabled {
+func Close() {
+	glog.Info("Closing auditor")
+	auditor.mu.Lock()
+	defer auditor.mu.Unlock()
+	auditor.tick.Stop()
+	auditor.log.Sync()
+}
+
+func (a *auditLogger) AuditFromEvent(event *AuditEvent) {
+	a.mu.RLock()
+	defer a.mu.RUnlock()
+	if a.log == nil {
 		return
 	}
-	b, _ := json.Marshal(event)
-	a.log.Print(string(b))
+	a.log.AuditI(event.Endpoint,
+		"user", event.User,
+		"server", event.ServerHost,
+		"client", event.ClientHost,
+		"req_body", event.Req,
+		"query_param", event.QueryParams,
+		"status", event.Status,
+		"time", event.TimeTaken)
 }
 
 type ResponseWriter struct {
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
new file mode 100644
index 00000000000..edd014f97e0
--- /dev/null
+++ b/ee/audit/interceptor.go
@@ -0,0 +1,103 @@
+package audit
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"github.com/dgraph-io/dgraph/x"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"time"
+)
+
+func AuditRequestGRPC(ctx context.Context, req interface{},
+	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+	startTime := time.Now().UnixNano()
+	response, err := handler(ctx, req)
+	maybeAuditGRPC(ctx, req, err, startTime)
+	return response, err
+}
+
+func AuditRequestHttp(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		startTime := time.Now().UnixNano()
+		lrw := NewResponseWriter(w)
+		var buf bytes.Buffer
+		tee := io.TeeReader(r.Body, &buf)
+		r.Body = ioutil.NopCloser(tee)
+		next.ServeHTTP(lrw, r)
+		r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
+		maybeAuditHttp(lrw, r, startTime)
+	})
+}
+
+func maybeAuditGRPC(ctx context.Context, req interface{}, err error, startTime int64) {
+	var code codes.Code
+	if serr, ok := status.FromError(err); !ok {
+		code = codes.Unknown
+	} else {
+		code = serr.Code()
+	}
+	clientHost := ""
+	p, ok := peer.FromContext(ctx)
+	if ok {
+		clientHost = p.Addr.String()
+	}
+
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		md = metadata.New(nil)
+	}
+
+	token := ""
+	t := md.Get("accessJwt")
+	if len(t) > 0 {
+		token = t[0]
+	}
+	event := &AuditEvent{
+		User:       getUserId(token),
+		ServerHost: x.WorkerConfig.MyAddr,
+		ClientHost: clientHost,
+		Req:        fmt.Sprintf("%v", req),
+		Status:     int(code),
+		TimeTaken:  time.Now().UnixNano() - startTime,
+	}
+	auditor.AuditFromEvent(event)
+}
+
+func maybeAuditHttp(w *ResponseWriter, r *http.Request, startTime int64) {
+	all, _ := ioutil.ReadAll(r.Body)
+	event := &AuditEvent{
+		User:        getUserId(r.Header.Get("X-Dgraph-AccessToken")),
+		ServerHost:  x.WorkerConfig.MyAddr,
+		ClientHost:  r.RemoteAddr,
+		Endpoint:    r.URL.Path,
+		Req:         string(all),
+		Status:      w.statusCode,
+		QueryParams: r.URL.Query(),
+		TimeTaken:   time.Now().UnixNano() - startTime,
+	}
+	auditor.AuditFromEvent(event)
+}
+
+func getUserId(token string) string {
+	var userId string
+	var err error
+	if token == "" {
+		if x.WorkerConfig.AclEnabled {
+			userId = UnauthorisedUser
+		}
+	} else {
+		userId, err = x.ExtractUserName(token)
+		if err != nil {
+			userId = UnknownUser
+		}
+	}
+	return userId
+}
diff --git a/posting/list.go b/posting/list.go
index 9b8101b645f..1d97b7d5aa5 100644
--- a/posting/list.go
+++ b/posting/list.go
@@ -31,7 +31,6 @@ import (
 	"github.com/dgraph-io/badger/v3/y"
 	"github.com/dgraph-io/dgraph/algo"
 	"github.com/dgraph-io/dgraph/codec"
-	"github.com/dgraph-io/dgraph/dgraph/cmd/zero"
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/schema"
 	"github.com/dgraph-io/dgraph/types"
@@ -499,7 +498,7 @@ func (l *List) addMutationInternal(ctx context.Context, txn *Txn, t *pb.Directed
 	l.AssertLock()
 
 	if txn.ShouldAbort() {
-		return zero.ErrConflict
+		return x.ErrConflict
 	}
 
 	mpost := NewPosting(t)
diff --git a/worker/draft.go b/worker/draft.go
index c508adf7e92..a1f8313a2dc 100644
--- a/worker/draft.go
+++ b/worker/draft.go
@@ -41,7 +41,6 @@ import (
 	"github.com/dgraph-io/badger/v3"
 	bpb "github.com/dgraph-io/badger/v3/pb"
 	"github.com/dgraph-io/dgraph/conn"
-	"github.com/dgraph-io/dgraph/dgraph/cmd/zero"
 	"github.com/dgraph-io/dgraph/posting"
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/raftwal"
@@ -490,7 +489,7 @@ func (n *node) applyMutations(ctx context.Context, proposal *pb.Proposal) (rerr
 	txn := posting.Oracle().RegisterStartTs(m.StartTs)
 	if txn.ShouldAbort() {
 		span.Annotatef(nil, "Txn %d should abort.", m.StartTs)
-		return zero.ErrConflict
+		return x.ErrConflict
 	}
 	// Discard the posting lists from cache to release memory at the end.
 	defer txn.Update()
diff --git a/x/logger.go b/x/logger.go
index 2ead08f21e9..cab19bd0163 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -1,41 +1,55 @@
 package x
 
 import (
-	"os"
-
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
-	"golang.org/x/net/context"
+	"gopkg.in/natefinch/lumberjack.v2"
 )
 
-var logging LoggerImpl
-
 type ILogger interface {
-	Audit(ctx context.Context, args ...interface{})
+	AuditI(string, ...interface{})
 	Sync()
 }
 
-func initLogger() error {
+func InitLogger(dir string, filename string) (ILogger, error) {
+	getWriterSyncer := func() zapcore.WriteSyncer {
+		return zapcore.AddSync(&lumberjack.Logger{
+			Filename: dir + "/" + filename,
+			MaxSize:  100,
+			MaxAge:   30,
+		})
+	}
 	core := zapcore.NewCore(zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
-		zapcore.AddSync(os.Stdout), zap.DebugLevel)
+		getWriterSyncer(), zap.DebugLevel)
 
-	logger := zap.New(core)
-	sugarLogger := logger.Sugar()
-
-	logging = LoggerImpl{
-		sugarLogger: sugarLogger,
-	}
-	return nil
+	return &LoggerImpl{
+		logger: zap.New(core),
+	}, nil
 }
 
 type LoggerImpl struct {
-	sugarLogger *zap.SugaredLogger
+	logger *zap.Logger
 }
 
-func Audit(args ...interface{}) {
-	logging.sugarLogger.Info(args...)
+func (l *LoggerImpl) AuditI(msg string, args ...interface{}) {
+	flds := make([]zap.Field, len(args)/2)
+	for i := 0; i < len(args); i = i + 2 {
+		flds[i/2] = zap.Any(args[i].(string), args[i+1])
+	}
+	l.logger.Info(msg, flds...)
+}
+
+func (l *LoggerImpl) AuditE(msg string, args ...interface{}) {
+	flds := make([]zap.Field, len(args)/2)
+	for i := range args {
+		flds[i/2] = zap.Any(args[i].(string), args[i+1])
+	}
+	l.logger.Error(msg, flds...)
 }
 
 func (l *LoggerImpl) Sync() {
-	_ = l.sugarLogger.Sync()
+	if l == nil {
+		return
+	}
+	_ = l.logger.Sync()
 }
diff --git a/x/x.go b/x/x.go
index cd9a7016c8f..baab8c1e11e 100644
--- a/x/x.go
+++ b/x/x.go
@@ -67,6 +67,8 @@ var (
 	ErrNoJwt        = errors.New("no accessJwt available")
 	// ErrorInvalidLogin is returned when username or password is incorrect in login
 	ErrorInvalidLogin = errors.New("invalid username or password")
+	// ErrConflict is returned when commit couldn't succeed due to conflicts.
+	ErrConflict = errors.New("Transaction conflict")
 )
 
 const (

From 6ba4f3b5432813a833040b062c0e4a9babff2173 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Thu, 14 Jan 2021 19:37:48 +0530
Subject: [PATCH 03/31] making audit logs to work with zero too

---
 dgraph/cmd/alpha/run.go | 10 ++-----
 dgraph/cmd/zero/run.go  | 22 ++++++++++----
 ee/audit/audit.go       | 65 +++++++++++++++++++----------------------
 ee/audit/interceptor.go | 15 ++++++++--
 worker/config.go        |  1 -
 x/config.go             |  3 --
 x/logger.go             | 23 ++++++++-------
 7 files changed, 75 insertions(+), 64 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 9ee15a4055c..aaacf8b1456 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -614,7 +614,6 @@ func run() {
 		MutationsMode: worker.AllowMutations,
 		AuthToken:     Alpha.Conf.GetString("auth_token"),
 		AuditEnabled:  Alpha.Conf.GetBool("audit_enabled"),
-		AuditDir:      Alpha.Conf.GetString("audit_dir"),
 	}
 
 	secretFile := Alpha.Conf.GetString("acl_secret_file")
@@ -676,7 +675,6 @@ func run() {
 		LudicrousConcurrency: Alpha.Conf.GetInt("ludicrous_concurrency"),
 		TLSClientConfig:      tlsClientConf,
 		TLSServerConfig:      tlsServerConf,
-		AuditEnabled:         Alpha.Conf.GetBool("audit_enabled"),
 		HmacSecret:           opts.HmacSecret,
 	}
 	x.WorkerConfig.Parse(Alpha.Conf)
@@ -800,10 +798,8 @@ func run() {
 	// close alpha. This closer is for closing and waiting that subscription.
 	adminCloser := z.NewCloser(1)
 
-	// Audit is enterprise feature. If enabled, audit logs will be generate in the audit directory.
-	if opts.AuditEnabled {
-		go audit.InitAuditorIfNecessary(opts.AuditDir)
-	}
+	// Audit is enterprise feature.
+	audit.InitAuditorIfNecessary(Alpha.Conf, worker.EnterpriseEnabled)
 
 	setupServer(adminCloser)
 	glog.Infoln("GRPC and HTTP stopped.")
@@ -819,7 +815,7 @@ func run() {
 	glog.Infoln("adminCloser closed.")
 
 	audit.Close()
-
+	glog.Infoln("audit logs are closed.")
 	worker.State.Dispose()
 	x.RemoveCidFile()
 	glog.Info("worker.State disposed.")
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index 9a76e23f875..390fd1b186e 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"github.com/dgraph-io/dgraph/ee/audit"
 	"log"
 	"net"
 	"net/http"
@@ -123,6 +124,7 @@ func (st *state) serveGRPC(l net.Listener, store *raftwal.DiskStorage) {
 		grpc.MaxSendMsgSize(x.GrpcMaxSize),
 		grpc.MaxConcurrentStreams(1000),
 		grpc.StatsHandler(&ocgrpc.ServerHandler{}),
+		grpc.UnaryInterceptor(audit.AuditRequestGRPC),
 	}
 
 	tlsConf, err := x.LoadServerTLSConfigForInternalPort(Zero.Conf)
@@ -260,17 +262,25 @@ func run() {
 	go x.StartListenHttpAndHttps(httpListener, tlsCfg, st.zero.closer)
 
 	http.HandleFunc("/health", st.pingResponse)
-	http.HandleFunc("/state", st.getState)
-	http.HandleFunc("/removeNode", st.removeNode)
-	http.HandleFunc("/moveTablet", st.moveTablet)
-	http.HandleFunc("/assign", st.assign)
-	http.HandleFunc("/enterpriseLicense", st.applyEnterpriseLicense)
+	http.Handle("/state", audit.AuditRequestHttp(http.HandlerFunc(st.getState)))
+	http.Handle("/removeNode", audit.AuditRequestHttp(http.HandlerFunc(st.removeNode)))
+	http.Handle("/moveTablet", audit.AuditRequestHttp(http.HandlerFunc(st.moveTablet)))
+	http.Handle("/assign", audit.AuditRequestHttp(http.HandlerFunc(st.assign)))
+	http.Handle("/enterpriseLicense",
+		audit.AuditRequestHttp(http.HandlerFunc(st.applyEnterpriseLicense)))
 	http.HandleFunc("/jemalloc", x.JemallocHandler)
 	zpages.Handle(http.DefaultServeMux, "/z")
 
 	// This must be here. It does not work if placed before Grpc init.
 	x.Check(st.node.initAndStartNode())
 
+	audit.InitAuditorIfNecessary(Zero.Conf, func() bool {
+		if st.zero == nil || st.zero.state == nil {
+			return false
+		}
+		return st.zero.state.GetLicense().GetEnabled()
+	})
+
 	if Zero.Conf.GetBool("telemetry") {
 		go st.zero.periodicallyPostTelemetry()
 	}
@@ -326,4 +336,6 @@ func run() {
 	glog.Infof("Raft WAL closed with err: %v\n", err)
 	st.zero.orc.close()
 	glog.Infoln("All done. Goodbye!")
+	audit.Close()
+	glog.Infoln("audit logs are closed")
 }
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index 3011160e26a..79b7bfd3ca3 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -1,14 +1,13 @@
 package audit
 
 import (
+	"github.com/spf13/viper"
 	"net/http"
-	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/golang/glog"
 
-	"github.com/dgraph-io/dgraph/worker"
 	"github.com/dgraph-io/dgraph/x"
 )
 
@@ -33,67 +32,63 @@ const (
 var auditor *auditLogger
 
 type auditLogger struct {
-	mu   sync.RWMutex
+	log  *x.Logger
 	tick *time.Ticker
-	log  x.ILogger
 }
 
-func InitAuditorIfNecessary(dir string) {
+func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
+	if !conf.GetBool("audit_enabled") {
+		return
+	}
 	auditor = &auditLogger{
-		mu: sync.RWMutex{},
-		tick: time.NewTicker(time.Minute * 5),
+		tick: time.NewTicker(time.Second * 5),
 	}
-
-	if !worker.EnterpriseEnabled() {
-		return
+	if eeEnabled() {
+		auditor.log = initlog(conf.GetString("audit_dir"))
+		glog.Infoln("audit logs are enabled")
 	}
 
-	initlog := func() x.ILogger{
-		logger, err := x.InitLogger(dir, "dgraph_audit.log")
-		if err != nil {
-			glog.Errorf("error while initiating auditor %v", err)
-			return nil
-		}
-		return logger
+	go checkIfEEEnabled(eeEnabled, conf.GetString("audit_dir"))
+}
+
+func initlog(dir string) *x.Logger {
+	logger, err := x.InitLogger(dir, "dgraph_audit.log")
+	if err != nil {
+		glog.Errorf("error while initiating auditor %v", err)
+		return nil
 	}
-	auditor.log = initlog()
+	return logger
+}
+
+func checkIfEEEnabled(eeEnabledFunc func() bool, dir string) {
 	for {
 		select {
-		case <- auditor.tick.C:
-			if !worker.EnterpriseEnabled() {
+		case <-auditor.tick.C:
+			if !eeEnabledFunc() {
 				if atomic.LoadUint32(&auditEnabled) != 0 {
+					glog.Infof("audit logs are disabled")
 					atomic.StoreUint32(&auditEnabled, 0)
-					auditor.mu.Lock()
+					auditor.log.Sync()
 					auditor.log = nil
-					auditor.mu.Unlock()
-					continue
 				}
+				continue
 			}
 
 			if atomic.LoadUint32(&auditEnabled) != 1 {
+				glog.Info("audit logs are enabled")
+				auditor.log = initlog(dir)
 				atomic.StoreUint32(&auditEnabled, 1)
-				auditor.mu.Lock()
-				auditor.log = initlog()
-				auditor.mu.Unlock()
 			}
 		}
 	}
 }
 
 func Close() {
-	glog.Info("Closing auditor")
-	auditor.mu.Lock()
-	defer auditor.mu.Unlock()
 	auditor.tick.Stop()
 	auditor.log.Sync()
 }
 
-func (a *auditLogger) AuditFromEvent(event *AuditEvent) {
-	a.mu.RLock()
-	defer a.mu.RUnlock()
-	if a.log == nil {
-		return
-	}
+func (a *auditLogger) AuditEvent(event *AuditEvent) {
 	a.log.AuditI(event.Endpoint,
 		"user", event.User,
 		"server", event.ServerHost,
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index edd014f97e0..3b34b787ed8 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -13,11 +13,16 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"sync/atomic"
 	"time"
 )
 
 func AuditRequestGRPC(ctx context.Context, req interface{},
 	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+	if atomic.LoadUint32(&auditEnabled) == 0 {
+		return handler(ctx, req)
+	}
+
 	startTime := time.Now().UnixNano()
 	response, err := handler(ctx, req)
 	maybeAuditGRPC(ctx, req, err, startTime)
@@ -25,6 +30,12 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 }
 
 func AuditRequestHttp(next http.Handler) http.Handler {
+	if atomic.LoadUint32(&auditEnabled) == 0 {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			next.ServeHTTP(w, r)
+		})
+	}
+
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		startTime := time.Now().UnixNano()
 		lrw := NewResponseWriter(w)
@@ -68,7 +79,7 @@ func maybeAuditGRPC(ctx context.Context, req interface{}, err error, startTime i
 		Status:     int(code),
 		TimeTaken:  time.Now().UnixNano() - startTime,
 	}
-	auditor.AuditFromEvent(event)
+	auditor.AuditEvent(event)
 }
 
 func maybeAuditHttp(w *ResponseWriter, r *http.Request, startTime int64) {
@@ -83,7 +94,7 @@ func maybeAuditHttp(w *ResponseWriter, r *http.Request, startTime int64) {
 		QueryParams: r.URL.Query(),
 		TimeTaken:   time.Now().UnixNano() - startTime,
 	}
-	auditor.AuditFromEvent(event)
+	auditor.AuditEvent(event)
 }
 
 func getUserId(token string) string {
diff --git a/worker/config.go b/worker/config.go
index e0b6076395a..9bbc1af4077 100644
--- a/worker/config.go
+++ b/worker/config.go
@@ -71,7 +71,6 @@ type Options struct {
 	// CacheMb is the total memory allocated between all the caches.
 	CacheMb int64
 
-	AuditDir     string
 	AuditEnabled bool
 }
 
diff --git a/x/config.go b/x/config.go
index b030056e7f5..607848ba078 100644
--- a/x/config.go
+++ b/x/config.go
@@ -112,9 +112,6 @@ type WorkerOptions struct {
 	LogRequest int32
 	// If true, we should call msync or fsync after every write to survive hard reboots.
 	HardSync bool
-
-	// AuditEnabled determines whether the audit feature is enabled or not
-	AuditEnabled bool
 }
 
 // WorkerConfig stores the global instance of the worker package's options.
diff --git a/x/logger.go b/x/logger.go
index cab19bd0163..6096e56f448 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -6,12 +6,7 @@ import (
 	"gopkg.in/natefinch/lumberjack.v2"
 )
 
-type ILogger interface {
-	AuditI(string, ...interface{})
-	Sync()
-}
-
-func InitLogger(dir string, filename string) (ILogger, error) {
+func InitLogger(dir string, filename string) (*Logger, error) {
 	getWriterSyncer := func() zapcore.WriteSyncer {
 		return zapcore.AddSync(&lumberjack.Logger{
 			Filename: dir + "/" + filename,
@@ -22,16 +17,19 @@ func InitLogger(dir string, filename string) (ILogger, error) {
 	core := zapcore.NewCore(zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
 		getWriterSyncer(), zap.DebugLevel)
 
-	return &LoggerImpl{
+	return &Logger{
 		logger: zap.New(core),
 	}, nil
 }
 
-type LoggerImpl struct {
+type Logger struct {
 	logger *zap.Logger
 }
 
-func (l *LoggerImpl) AuditI(msg string, args ...interface{}) {
+func (l *Logger) AuditI(msg string, args ...interface{}) {
+	if l == nil {
+		return
+	}
 	flds := make([]zap.Field, len(args)/2)
 	for i := 0; i < len(args); i = i + 2 {
 		flds[i/2] = zap.Any(args[i].(string), args[i+1])
@@ -39,7 +37,10 @@ func (l *LoggerImpl) AuditI(msg string, args ...interface{}) {
 	l.logger.Info(msg, flds...)
 }
 
-func (l *LoggerImpl) AuditE(msg string, args ...interface{}) {
+func (l *Logger) AuditE(msg string, args ...interface{}) {
+	if l == nil {
+		return
+	}
 	flds := make([]zap.Field, len(args)/2)
 	for i := range args {
 		flds[i/2] = zap.Any(args[i].(string), args[i+1])
@@ -47,7 +48,7 @@ func (l *LoggerImpl) AuditE(msg string, args ...interface{}) {
 	l.logger.Error(msg, flds...)
 }
 
-func (l *LoggerImpl) Sync() {
+func (l *Logger) Sync() {
 	if l == nil {
 		return
 	}

From 8f4b3bc0d5fba6c19621eed18149166a70ab43af Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Thu, 14 Jan 2021 19:42:30 +0530
Subject: [PATCH 04/31] go mod tidy

---
 go.mod | 1 -
 go.sum | 7 +++----
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/go.mod b/go.mod
index e46498fafeb..b77791e7fe5 100644
--- a/go.mod
+++ b/go.mod
@@ -60,7 +60,6 @@ require (
 	github.com/twpayne/go-geom v1.0.5
 	go.etcd.io/etcd v0.0.0-20190228193606-a943ad0ee4c9
 	go.opencensus.io v0.22.5
-	go.uber.org/multierr v1.6.0 // indirect
 	go.uber.org/zap v1.16.0
 	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9
 	golang.org/x/net v0.0.0-20201021035429-f5854403a974
diff --git a/go.sum b/go.sum
index 0e1c63af36f..d300de1abbf 100644
--- a/go.sum
+++ b/go.sum
@@ -590,14 +590,11 @@ go.uber.org/atomic v1.4.0 h1:cxzIVoETapQEqDhQu3QfnvXAV4AlzcvUCxkVUFw3+EU=
 go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
 go.uber.org/atomic v1.6.0 h1:Ezj3JGmsOnG1MoRWQkPBsKLe9DwWD9QeXzTRzzldNVk=
 go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
-go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
-go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/multierr v1.1.0 h1:HoEmRHQPVSqub6w2z2d2EOVs2fjyFRGyofhKuyDq0QI=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/multierr v1.5.0 h1:KCa4XfM8CWFCpxXRGok+Q0SS/0XBhMDbHHGABQLvD2A=
 go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU=
-go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
-go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
+go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee h1:0mgffUl7nfd+FpvXMVz4IDEaUSmT1ysygQC7qYo7sG4=
 go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA=
 go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
 go.uber.org/zap v1.10.0 h1:ORx85nbTijNz8ljznvCMR1ZBIPKFn3jQrag10X2AsuM=
@@ -632,6 +629,7 @@ golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHl
 golang.org/x/lint v0.0.0-20190409202823-959b441ac422/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/lint v0.0.0-20190909230951-414d861bb4ac/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
+golang.org/x/lint v0.0.0-20200302205851-738671d3881b h1:Wh+f8QHJXR411sJR8/vRBTZ7YapZaRvUcLFFJhusH0k=
 golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
 golang.org/x/mobile v0.0.0-20190312151609-d3739f865fa6/go.mod h1:z+o9i4GpDbdi3rU15maQ/Ox0txvL9dWGYEHz965HBQE=
 golang.org/x/mobile v0.0.0-20190719004257-d2bd2a29d028/go.mod h1:E/iHnbuqvinMTCcRqshq8CkpyQDoeVncDDYHnLhea+o=
@@ -823,6 +821,7 @@ honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
+honnef.co/go/tools v0.0.1-2019.2.3 h1:3JgtbtFHMiCmsznwGVTUWbgGov+pVqnlf1dEJTNAXeM=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=

From 33aff9cd4bc90d7dbfc385c8556d47f6fdec5243 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Thu, 14 Jan 2021 19:53:06 +0530
Subject: [PATCH 05/31] fixing minor things to clean

---
 ee/audit/audit.go | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index 79b7bfd3ca3..c747873b28a 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -45,10 +45,11 @@ func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
 	}
 	if eeEnabled() {
 		auditor.log = initlog(conf.GetString("audit_dir"))
+		atomic.StoreUint32(&auditEnabled, 1)
 		glog.Infoln("audit logs are enabled")
 	}
 
-	go checkIfEEEnabled(eeEnabled, conf.GetString("audit_dir"))
+	go trackIfEEValid(eeEnabled, conf.GetString("audit_dir"))
 }
 
 func initlog(dir string) *x.Logger {
@@ -60,14 +61,14 @@ func initlog(dir string) *x.Logger {
 	return logger
 }
 
-func checkIfEEEnabled(eeEnabledFunc func() bool, dir string) {
+func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
 	for {
 		select {
 		case <-auditor.tick.C:
 			if !eeEnabledFunc() {
 				if atomic.LoadUint32(&auditEnabled) != 0 {
-					glog.Infof("audit logs are disabled")
 					atomic.StoreUint32(&auditEnabled, 0)
+					glog.Infof("audit logs are disabled")
 					auditor.log.Sync()
 					auditor.log = nil
 				}
@@ -75,9 +76,9 @@ func checkIfEEEnabled(eeEnabledFunc func() bool, dir string) {
 			}
 
 			if atomic.LoadUint32(&auditEnabled) != 1 {
-				glog.Info("audit logs are enabled")
 				auditor.log = initlog(dir)
 				atomic.StoreUint32(&auditEnabled, 1)
+				glog.Info("audit logs are enabled")
 			}
 		}
 	}

From d78651fba0ea4b4728d03064864638230f0bb975 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Thu, 14 Jan 2021 20:02:38 +0530
Subject: [PATCH 06/31] adding endpoint to grpc audit

---
 ee/audit/interceptor.go | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 3b34b787ed8..9a6ffa889c5 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -25,7 +25,7 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 
 	startTime := time.Now().UnixNano()
 	response, err := handler(ctx, req)
-	maybeAuditGRPC(ctx, req, err, startTime)
+	maybeAuditGRPC(ctx, req, info, err, startTime)
 	return response, err
 }
 
@@ -48,7 +48,8 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 	})
 }
 
-func maybeAuditGRPC(ctx context.Context, req interface{}, err error, startTime int64) {
+func maybeAuditGRPC(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error,
+	startTime int64) {
 	var code codes.Code
 	if serr, ok := status.FromError(err); !ok {
 		code = codes.Unknown
@@ -75,6 +76,7 @@ func maybeAuditGRPC(ctx context.Context, req interface{}, err error, startTime i
 		User:       getUserId(token),
 		ServerHost: x.WorkerConfig.MyAddr,
 		ClientHost: clientHost,
+		Endpoint:   info.FullMethod,
 		Req:        fmt.Sprintf("%v", req),
 		Status:     int(code),
 		TimeTaken:  time.Now().UnixNano() - startTime,

From 8f33f61f13f5b956103f68a75b3cb7e5a9e9c08e Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Thu, 14 Jan 2021 21:32:39 +0530
Subject: [PATCH 07/31] fixing alpha and zero default directories

---
 dgraph/cmd/alpha/run.go |  3 +--
 dgraph/cmd/zero/run.go  |  9 ++++-----
 ee/audit/audit.go       |  2 +-
 ee/audit/interceptor.go | 10 ++++------
 x/logger.go             | 12 +++++++++++-
 5 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index aaacf8b1456..2b930a96c0f 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -199,8 +199,7 @@ they form a Raft group and provide synchronous replication.
 		PostingListCache,PstoreBlockCache,PstoreIndexCache,WAL).`)
 
 	flag.Bool("audit_enabled", false, "Set to true to enable audit logs.")
-	// todo(aman): check what to set audit_dir default to
-	flag.String("audit_dir", "./", "Set path to directory where to save the audit logs.")
+	flag.String("audit_dir", "aa", "Set path to directory where to save the audit logs.")
 
 	// TLS configurations
 	x.RegisterServerTLSFlags(flag)
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index 390fd1b186e..f7e31b5ce95 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -98,8 +98,7 @@ instances to achieve high-availability.
 	flag.String("enterprise_license", "", "Path to the enterprise license file.")
 
 	flag.Bool("audit_enabled", false, "Set to true to enable audit logs.")
-	// todo(aman): check what to set audit_dir default to
-	flag.String("audit_dir", "./", "Set path to directory where to save the audit logs.")
+	flag.String("audit_dir", "za", "Set path to directory where to save the audit logs.")
 
 	// TLS configurations
 	x.RegisterServerTLSFlags(flag)
@@ -271,9 +270,6 @@ func run() {
 	http.HandleFunc("/jemalloc", x.JemallocHandler)
 	zpages.Handle(http.DefaultServeMux, "/z")
 
-	// This must be here. It does not work if placed before Grpc init.
-	x.Check(st.node.initAndStartNode())
-
 	audit.InitAuditorIfNecessary(Zero.Conf, func() bool {
 		if st.zero == nil || st.zero.state == nil {
 			return false
@@ -281,6 +277,9 @@ func run() {
 		return st.zero.state.GetLicense().GetEnabled()
 	})
 
+	// This must be here. It does not work if placed before Grpc init.
+	x.Check(st.node.initAndStartNode())
+
 	if Zero.Conf.GetBool("telemetry") {
 		go st.zero.periodicallyPostTelemetry()
 	}
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index c747873b28a..9ffa2c18e20 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -78,7 +78,7 @@ func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
 			if atomic.LoadUint32(&auditEnabled) != 1 {
 				auditor.log = initlog(dir)
 				atomic.StoreUint32(&auditEnabled, 1)
-				glog.Info("audit logs are enabled")
+				glog.Infof("audit logs are enabled")
 			}
 		}
 	}
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 9a6ffa889c5..1ed1215f3c7 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -30,13 +30,11 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 }
 
 func AuditRequestHttp(next http.Handler) http.Handler {
-	if atomic.LoadUint32(&auditEnabled) == 0 {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			next.ServeHTTP(w, r)
-		})
-	}
-
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if atomic.LoadUint32(&auditEnabled) == 0 {
+			next.ServeHTTP(w, r)
+			return
+		}
 		startTime := time.Now().UnixNano()
 		lrw := NewResponseWriter(w)
 		var buf bytes.Buffer
diff --git a/x/logger.go b/x/logger.go
index 6096e56f448..932c3603c25 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -4,12 +4,22 @@ import (
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	"gopkg.in/natefinch/lumberjack.v2"
+	"os"
+	"path/filepath"
 )
 
 func InitLogger(dir string, filename string) (*Logger, error) {
+	err := os.MkdirAll(dir, 0700)
+	if err != nil {
+		return nil, err
+	}
+	path, err := filepath.Abs(filepath.Join(dir, filename))
+	if err != nil {
+		return nil, err
+	}
 	getWriterSyncer := func() zapcore.WriteSyncer {
 		return zapcore.AddSync(&lumberjack.Logger{
-			Filename: dir + "/" + filename,
+			Filename: path,
 			MaxSize:  100,
 			MaxAge:   30,
 		})

From eb2b7c7c55a168c790c20f1f9ae4461fd8893e66 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Fri, 15 Jan 2021 11:01:09 +0530
Subject: [PATCH 08/31] fixing alpha and zero default directories

---
 dgraph/cmd/alpha/run.go |  1 +
 dgraph/cmd/zero/run.go  |  9 +++++++++
 ee/audit/audit.go       | 19 +++----------------
 ee/audit/interceptor.go | 16 ++++++++++++++++
 worker/config.go        |  5 +++++
 x/config.go             |  2 ++
 6 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 2b930a96c0f..73253040c29 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -675,6 +675,7 @@ func run() {
 		TLSClientConfig:      tlsClientConf,
 		TLSServerConfig:      tlsServerConf,
 		HmacSecret:           opts.HmacSecret,
+		AuditDir:             Alpha.Conf.GetString("audit_dir"),
 	}
 	x.WorkerConfig.Parse(Alpha.Conf)
 	x.CheckFlag(x.WorkerConfig.Raft, "group", "idx", "learner")
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index f7e31b5ce95..f7cfaa47c09 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -26,6 +26,7 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"path/filepath"
 	"syscall"
 	"time"
 
@@ -220,6 +221,14 @@ func run() {
 		}
 	}
 
+	wd, err := filepath.Abs(opts.w)
+	x.Check(err)
+	ad, err := filepath.Abs(Zero.Conf.GetString("audit_dir"))
+	x.Check(err)
+	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').",
+		Zero.Conf.Get("audit_dir"))
+
+
 	if opts.rebalanceInterval <= 0 {
 		log.Fatalf("ERROR: Rebalance interval must be greater than zero. Found: %d",
 			opts.rebalanceInterval)
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index 9ffa2c18e20..ca860650800 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -2,7 +2,6 @@ package audit
 
 import (
 	"github.com/spf13/viper"
-	"net/http"
 	"sync/atomic"
 	"time"
 
@@ -85,6 +84,9 @@ func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
 }
 
 func Close() {
+	if auditor == nil {
+		return
+	}
 	auditor.tick.Stop()
 	auditor.log.Sync()
 }
@@ -100,18 +102,3 @@ func (a *auditLogger) AuditEvent(event *AuditEvent) {
 		"time", event.TimeTaken)
 }
 
-type ResponseWriter struct {
-	http.ResponseWriter
-	statusCode int
-}
-
-func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
-	// WriteHeader(int) is not called if our response implicitly returns 200 OK, so
-	// we default to that status code.
-	return &ResponseWriter{w, http.StatusOK}
-}
-
-func (lrw *ResponseWriter) WriteHeader(code int) {
-	lrw.statusCode = code
-	lrw.ResponseWriter.WriteHeader(code)
-}
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 1ed1215f3c7..89053094594 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -112,3 +112,19 @@ func getUserId(token string) string {
 	}
 	return userId
 }
+
+type ResponseWriter struct {
+	http.ResponseWriter
+	statusCode int
+}
+
+func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
+	// WriteHeader(int) is not called if our response implicitly returns 200 OK, so
+	// we default to that status code.
+	return &ResponseWriter{w, http.StatusOK}
+}
+
+func (lrw *ResponseWriter) WriteHeader(code int) {
+	lrw.statusCode = code
+	lrw.ResponseWriter.WriteHeader(code)
+}
\ No newline at end of file
diff --git a/worker/config.go b/worker/config.go
index 9bbc1af4077..d9256459c31 100644
--- a/worker/config.go
+++ b/worker/config.go
@@ -96,7 +96,12 @@ func (opt *Options) validate() {
 	x.Check(err)
 	td, err := filepath.Abs(x.WorkerConfig.TmpDir)
 	x.Check(err)
+	ad, err := filepath.Abs(x.WorkerConfig.AuditDir)
+	x.Check(err)
 	x.AssertTruef(pd != wd, "Posting and WAL directory cannot be the same ('%s').", opt.PostingDir)
 	x.AssertTruef(pd != td, "Posting and Tmp directory cannot be the same ('%s').", opt.PostingDir)
 	x.AssertTruef(wd != td, "WAL and Tmp directory cannot be the same ('%s').", opt.WALDir)
+	x.AssertTruef(ad != pd, "Posting and Audit Directory cannot be the same ('%s').", x.WorkerConfig.AuditDir)
+	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", x.WorkerConfig.AuditDir)
+	x.AssertTruef(ad != td, "Tmp and Audit directory cannot be the same ('%s').", x.WorkerConfig.AuditDir)
 }
diff --git a/x/config.go b/x/config.go
index 607848ba078..2808c57a135 100644
--- a/x/config.go
+++ b/x/config.go
@@ -112,6 +112,8 @@ type WorkerOptions struct {
 	LogRequest int32
 	// If true, we should call msync or fsync after every write to survive hard reboots.
 	HardSync bool
+
+	AuditDir string
 }
 
 // WorkerConfig stores the global instance of the worker package's options.

From 7515a1bb8fcc15365ea0ed95c3835c1bd53801cb Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Fri, 15 Jan 2021 12:08:26 +0530
Subject: [PATCH 09/31] adding skip method for grpc audits

---
 dgraph/cmd/zero/run.go  |  4 ++--
 ee/audit/audit.go       | 24 ++++++++++++++----------
 ee/audit/interceptor.go | 32 +++++++++++++++++++++++++-------
 3 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index f7cfaa47c09..a7ab5819f59 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -20,7 +20,6 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
-	"github.com/dgraph-io/dgraph/ee/audit"
 	"log"
 	"net"
 	"net/http"
@@ -30,6 +29,8 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/dgraph-io/dgraph/ee/audit"
+
 	"go.opencensus.io/plugin/ocgrpc"
 	otrace "go.opencensus.io/trace"
 	"go.opencensus.io/zpages"
@@ -228,7 +229,6 @@ func run() {
 	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').",
 		Zero.Conf.Get("audit_dir"))
 
-
 	if opts.rebalanceInterval <= 0 {
 		log.Fatalf("ERROR: Rebalance interval must be greater than zero. Found: %d",
 			opts.rebalanceInterval)
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index ca860650800..f407b8f6d95 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -1,10 +1,11 @@
 package audit
 
 import (
-	"github.com/spf13/viper"
 	"sync/atomic"
 	"time"
 
+	"github.com/spf13/viper"
+
 	"github.com/golang/glog"
 
 	"github.com/dgraph-io/dgraph/x"
@@ -13,19 +14,22 @@ import (
 var auditEnabled uint32
 
 type AuditEvent struct {
-	User        string              `json:"user"`
-	ServerHost  string              `json:"server_host"`
-	ClientHost  string              `json:"client_host"`
-	Endpoint    string              `json:"endpoint"`
-	Req         string              `json:"req"`
-	Status      int                 `json:"status"`
-	QueryParams map[string][]string `json:"query_params"`
-	TimeTaken   int64               `json:"time_taken"`
+	User        string
+	ServerHost  string
+	ClientHost  string
+	Endpoint    string
+	ReqType     string
+	Req         string
+	Status      int
+	QueryParams map[string][]string
+	TimeTaken   int64
 }
 
 const (
 	UnauthorisedUser = "UnauthorisedUser"
 	UnknownUser      = "UnknownUser"
+	Grpc             = "Grpc"
+	Http             = "Http"
 )
 
 var auditor *auditLogger
@@ -96,9 +100,9 @@ func (a *auditLogger) AuditEvent(event *AuditEvent) {
 		"user", event.User,
 		"server", event.ServerHost,
 		"client", event.ClientHost,
+		"req_type", event.ReqType,
 		"req_body", event.Req,
 		"query_param", event.QueryParams,
 		"status", event.Status,
 		"time", event.TimeTaken)
 }
-
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 89053094594..f510d07b373 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -4,22 +4,38 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"strings"
+	"sync/atomic"
+	"time"
+
 	"github.com/dgraph-io/dgraph/x"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
-	"io"
-	"io/ioutil"
-	"net/http"
-	"sync/atomic"
-	"time"
 )
 
 func AuditRequestGRPC(ctx context.Context, req interface{},
 	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-	if atomic.LoadUint32(&auditEnabled) == 0 {
+	skip := func(method string) bool {
+		skipApis := []string{"Heartbeat", "RaftMessage", "JoinCluster", "IsPeer", // raft server
+			"StreamMembership", "UpdateMembership", "Oracle", // zero server
+			"Check", "Watch", // health server
+		}
+
+		for _, api := range skipApis {
+			if strings.HasSuffix(method, api) {
+				return true
+			}
+		}
+		return false
+	}
+
+	if atomic.LoadUint32(&auditEnabled) == 0 || skip(info.FullMethod) {
 		return handler(ctx, req)
 	}
 
@@ -75,6 +91,7 @@ func maybeAuditGRPC(ctx context.Context, req interface{}, info *grpc.UnaryServer
 		ServerHost: x.WorkerConfig.MyAddr,
 		ClientHost: clientHost,
 		Endpoint:   info.FullMethod,
+		ReqType:    Grpc,
 		Req:        fmt.Sprintf("%v", req),
 		Status:     int(code),
 		TimeTaken:  time.Now().UnixNano() - startTime,
@@ -89,6 +106,7 @@ func maybeAuditHttp(w *ResponseWriter, r *http.Request, startTime int64) {
 		ServerHost:  x.WorkerConfig.MyAddr,
 		ClientHost:  r.RemoteAddr,
 		Endpoint:    r.URL.Path,
+		ReqType:     Http,
 		Req:         string(all),
 		Status:      w.statusCode,
 		QueryParams: r.URL.Query(),
@@ -127,4 +145,4 @@ func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
 func (lrw *ResponseWriter) WriteHeader(code int) {
 	lrw.statusCode = code
 	lrw.ResponseWriter.WriteHeader(code)
-}
\ No newline at end of file
+}

From 6e3d7704febc6e15a821735a0be11cfd7adb5ad0 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Fri, 15 Jan 2021 12:11:46 +0530
Subject: [PATCH 10/31] renaming rw

---
 ee/audit/interceptor.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index f510d07b373..23339e428f4 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -142,7 +142,7 @@ func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
 	return &ResponseWriter{w, http.StatusOK}
 }
 
-func (lrw *ResponseWriter) WriteHeader(code int) {
-	lrw.statusCode = code
-	lrw.ResponseWriter.WriteHeader(code)
+func (rw *ResponseWriter) WriteHeader(code int) {
+	rw.statusCode = code
+	rw.ResponseWriter.WriteHeader(code)
 }

From 6706ac5022f9684fae676a3bde57b2e826728a4a Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Fri, 15 Jan 2021 14:29:44 +0530
Subject: [PATCH 11/31] refactoring basic things

---
 dgraph/cmd/zero/run.go  |  2 +-
 ee/audit/audit.go       | 22 +++++++++---------
 ee/audit/interceptor.go | 49 ++++++++++++++++++-----------------------
 3 files changed, 32 insertions(+), 41 deletions(-)

diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index a7ab5819f59..cb12045260e 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -280,7 +280,7 @@ func run() {
 	zpages.Handle(http.DefaultServeMux, "/z")
 
 	audit.InitAuditorIfNecessary(Zero.Conf, func() bool {
-		if st.zero == nil || st.zero.state == nil {
+		if st.zero.state == nil {
 			return false
 		}
 		return st.zero.state.GetLicense().GetEnabled()
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index f407b8f6d95..b40a41f67df 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -20,9 +20,8 @@ type AuditEvent struct {
 	Endpoint    string
 	ReqType     string
 	Req         string
-	Status      int
+	Status      string
 	QueryParams map[string][]string
-	TimeTaken   int64
 }
 
 const (
@@ -44,7 +43,7 @@ func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
 		return
 	}
 	auditor = &auditLogger{
-		tick: time.NewTicker(time.Second * 5),
+		tick: time.NewTicker(time.Minute * 5),
 	}
 	if eeEnabled() {
 		auditor.log = initlog(conf.GetString("audit_dir"))
@@ -64,17 +63,17 @@ func initlog(dir string) *x.Logger {
 	return logger
 }
 
+// trackIfEEValid tracks enterprise license of the cluster.
+// Right now alpha doesn't know about the enterprise/licence.
+// That's why we needed to track if the current node is part of enterprise edition cluster
 func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
 	for {
 		select {
 		case <-auditor.tick.C:
-			if !eeEnabledFunc() {
-				if atomic.LoadUint32(&auditEnabled) != 0 {
-					atomic.StoreUint32(&auditEnabled, 0)
-					glog.Infof("audit logs are disabled")
-					auditor.log.Sync()
-					auditor.log = nil
-				}
+			if !eeEnabledFunc() && atomic.CompareAndSwapUint32(&auditEnabled, 1, 0) {
+				glog.Infof("audit logs are disabled")
+				auditor.log.Sync()
+				auditor.log = nil
 				continue
 			}
 
@@ -103,6 +102,5 @@ func (a *auditLogger) AuditEvent(event *AuditEvent) {
 		"req_type", event.ReqType,
 		"req_body", event.Req,
 		"query_param", event.QueryParams,
-		"status", event.Status,
-		"time", event.TimeTaken)
+		"status", event.Status)
 }
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 23339e428f4..e2558ded0b7 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -9,7 +9,6 @@ import (
 	"net/http"
 	"strings"
 	"sync/atomic"
-	"time"
 
 	"github.com/dgraph-io/dgraph/x"
 	"google.golang.org/grpc"
@@ -39,9 +38,8 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 		return handler(ctx, req)
 	}
 
-	startTime := time.Now().UnixNano()
 	response, err := handler(ctx, req)
-	maybeAuditGRPC(ctx, req, info, err, startTime)
+	auditGrpc(ctx, req, info, err)
 	return response, err
 }
 
@@ -51,19 +49,18 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 			next.ServeHTTP(w, r)
 			return
 		}
-		startTime := time.Now().UnixNano()
-		lrw := NewResponseWriter(w)
+
+		rw := NewResponseWriter(w)
 		var buf bytes.Buffer
 		tee := io.TeeReader(r.Body, &buf)
 		r.Body = ioutil.NopCloser(tee)
-		next.ServeHTTP(lrw, r)
+		next.ServeHTTP(rw, r)
 		r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
-		maybeAuditHttp(lrw, r, startTime)
+		auditHttp(rw, r)
 	})
 }
 
-func maybeAuditGRPC(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error,
-	startTime int64) {
+func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error) {
 	var code codes.Code
 	if serr, ok := status.FromError(err); !ok {
 		code = codes.Unknown
@@ -71,21 +68,17 @@ func maybeAuditGRPC(ctx context.Context, req interface{}, info *grpc.UnaryServer
 		code = serr.Code()
 	}
 	clientHost := ""
-	p, ok := peer.FromContext(ctx)
-	if ok {
+	if p, ok := peer.FromContext(ctx); ok {
 		clientHost = p.Addr.String()
 	}
 
-	md, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		md = metadata.New(nil)
-	}
-
 	token := ""
-	t := md.Get("accessJwt")
-	if len(t) > 0 {
-		token = t[0]
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		if t := md.Get("accessJwt"); len(t) > 0 {
+			token = t[0]
+		}
 	}
+
 	event := &AuditEvent{
 		User:       getUserId(token),
 		ServerHost: x.WorkerConfig.MyAddr,
@@ -93,24 +86,25 @@ func maybeAuditGRPC(ctx context.Context, req interface{}, info *grpc.UnaryServer
 		Endpoint:   info.FullMethod,
 		ReqType:    Grpc,
 		Req:        fmt.Sprintf("%v", req),
-		Status:     int(code),
-		TimeTaken:  time.Now().UnixNano() - startTime,
+		Status:     code.String(),
 	}
 	auditor.AuditEvent(event)
 }
 
-func maybeAuditHttp(w *ResponseWriter, r *http.Request, startTime int64) {
-	all, _ := ioutil.ReadAll(r.Body)
+func auditHttp(w *ResponseWriter, r *http.Request) {
+	rb, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		rb = []byte(err.Error())
+	}
 	event := &AuditEvent{
 		User:        getUserId(r.Header.Get("X-Dgraph-AccessToken")),
 		ServerHost:  x.WorkerConfig.MyAddr,
 		ClientHost:  r.RemoteAddr,
 		Endpoint:    r.URL.Path,
 		ReqType:     Http,
-		Req:         string(all),
-		Status:      w.statusCode,
+		Req:         string(rb),
+		Status:      http.StatusText(w.statusCode),
 		QueryParams: r.URL.Query(),
-		TimeTaken:   time.Now().UnixNano() - startTime,
 	}
 	auditor.AuditEvent(event)
 }
@@ -123,8 +117,7 @@ func getUserId(token string) string {
 			userId = UnauthorisedUser
 		}
 	} else {
-		userId, err = x.ExtractUserName(token)
-		if err != nil {
+		if userId, err = x.ExtractUserName(token); err != nil {
 			userId = UnknownUser
 		}
 	}

From 72db97670d06bc8c45bc1c91c4e24716216cdd64 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Fri, 15 Jan 2021 14:54:19 +0530
Subject: [PATCH 12/31] audit fix

---
 ee/audit/interceptor.go | 11 ++++-------
 x/logger.go             | 19 ++++++++++---------
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index e2558ded0b7..55ef002bcd4 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -37,7 +37,6 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 	if atomic.LoadUint32(&auditEnabled) == 0 || skip(info.FullMethod) {
 		return handler(ctx, req)
 	}
-
 	response, err := handler(ctx, req)
 	auditGrpc(ctx, req, info, err)
 	return response, err
@@ -79,7 +78,7 @@ func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
 		}
 	}
 
-	event := &AuditEvent{
+	auditor.AuditEvent(&AuditEvent{
 		User:       getUserId(token),
 		ServerHost: x.WorkerConfig.MyAddr,
 		ClientHost: clientHost,
@@ -87,8 +86,7 @@ func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
 		ReqType:    Grpc,
 		Req:        fmt.Sprintf("%v", req),
 		Status:     code.String(),
-	}
-	auditor.AuditEvent(event)
+	})
 }
 
 func auditHttp(w *ResponseWriter, r *http.Request) {
@@ -96,7 +94,7 @@ func auditHttp(w *ResponseWriter, r *http.Request) {
 	if err != nil {
 		rb = []byte(err.Error())
 	}
-	event := &AuditEvent{
+	auditor.AuditEvent(&AuditEvent{
 		User:        getUserId(r.Header.Get("X-Dgraph-AccessToken")),
 		ServerHost:  x.WorkerConfig.MyAddr,
 		ClientHost:  r.RemoteAddr,
@@ -105,8 +103,7 @@ func auditHttp(w *ResponseWriter, r *http.Request) {
 		Req:         string(rb),
 		Status:      http.StatusText(w.statusCode),
 		QueryParams: r.URL.Query(),
-	}
-	auditor.AuditEvent(event)
+	})
 }
 
 func getUserId(token string) string {
diff --git a/x/logger.go b/x/logger.go
index 932c3603c25..faaea4c3fb5 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -1,16 +1,16 @@
 package x
 
 import (
+	"os"
+	"path/filepath"
+
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	"gopkg.in/natefinch/lumberjack.v2"
-	"os"
-	"path/filepath"
 )
 
 func InitLogger(dir string, filename string) (*Logger, error) {
-	err := os.MkdirAll(dir, 0700)
-	if err != nil {
+	if err := os.MkdirAll(dir, 0700); err != nil {
 		return nil, err
 	}
 	path, err := filepath.Abs(filepath.Join(dir, filename))
@@ -36,13 +36,14 @@ type Logger struct {
 	logger *zap.Logger
 }
 
+// AuditI logs audit message as info. args are key value pairs with key as string value
 func (l *Logger) AuditI(msg string, args ...interface{}) {
 	if l == nil {
 		return
 	}
-	flds := make([]zap.Field, len(args)/2)
+	flds := make([]zap.Field, 0)
 	for i := 0; i < len(args); i = i + 2 {
-		flds[i/2] = zap.Any(args[i].(string), args[i+1])
+		flds = append(flds, zap.Any(args[i].(string), args[i+1]))
 	}
 	l.logger.Info(msg, flds...)
 }
@@ -51,9 +52,9 @@ func (l *Logger) AuditE(msg string, args ...interface{}) {
 	if l == nil {
 		return
 	}
-	flds := make([]zap.Field, len(args)/2)
-	for i := range args {
-		flds[i/2] = zap.Any(args[i].(string), args[i+1])
+	flds := make([]zap.Field, 0)
+	for i := 0; i < len(args); i = i + 2 {
+		flds = append(flds, zap.Any(args[i].(string), args[i+1]))
 	}
 	l.logger.Error(msg, flds...)
 }

From 97729e5b427b3aa4626a2c97e2dfbc114b09b0d3 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Fri, 15 Jan 2021 20:02:41 +0530
Subject: [PATCH 13/31] making zero and alpha logs from the start itself.

---
 dgraph/cmd/alpha/run.go       |  2 +-
 dgraph/cmd/zero/license_ee.go |  3 +++
 dgraph/cmd/zero/raft.go       |  4 ++++
 dgraph/cmd/zero/run.go        | 18 +++++++-----------
 ee/audit/audit.go             | 36 ++++++++++++++++++++++++-----------
 ee/audit/interceptor.go       | 36 +++++++++++++++++++++++------------
 6 files changed, 64 insertions(+), 35 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 73253040c29..caf78e76022 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -815,7 +815,7 @@ func run() {
 	glog.Infoln("adminCloser closed.")
 
 	audit.Close()
-	glog.Infoln("audit logs are closed.")
+	glog.Infoln("audit logs if enabled are closed.")
 	worker.State.Dispose()
 	x.RemoveCidFile()
 	glog.Info("worker.State disposed.")
diff --git a/dgraph/cmd/zero/license_ee.go b/dgraph/cmd/zero/license_ee.go
index 95f8ef3c3d9..e6e7f7c99ce 100644
--- a/dgraph/cmd/zero/license_ee.go
+++ b/dgraph/cmd/zero/license_ee.go
@@ -15,6 +15,7 @@ package zero
 import (
 	"bytes"
 	"context"
+	"github.com/dgraph-io/dgraph/ee/audit"
 	"io/ioutil"
 	"math"
 	"net/http"
@@ -91,6 +92,8 @@ func (n *node) updateEnterpriseState(closer *z.Closer) {
 			active := time.Now().UTC().Before(expiry)
 			if !active {
 				n.server.expireLicense()
+				audit.Close()
+				glog.Infoln("audit logs if enabled are closed.")
 				glog.Warningf("Your enterprise license has expired and enterprise features are " +
 					"disabled. To continue using enterprise features, apply a valid license. To receive " +
 					"a new license, contact us at https://dgraph.io/contact.")
diff --git a/dgraph/cmd/zero/raft.go b/dgraph/cmd/zero/raft.go
index d91961190e0..3a8930880f4 100644
--- a/dgraph/cmd/zero/raft.go
+++ b/dgraph/cmd/zero/raft.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/binary"
 	"fmt"
+	"github.com/dgraph-io/dgraph/ee/audit"
 	"log"
 	"math"
 	"sort"
@@ -401,6 +402,9 @@ func (n *node) applyProposal(e raftpb.Entry) (uint64, error) {
 		// Check expiry and set enabled accordingly.
 		expiry := time.Unix(state.License.ExpiryTs, 0).UTC()
 		state.License.Enabled = time.Now().UTC().Before(expiry)
+		if state.License.Enabled {
+			audit.InitAuditor(opts.auditEnabled, opts.auditDir)
+		}
 	}
 	if p.Snapshot != nil {
 		if err := n.applySnapshot(p.Snapshot); err != nil {
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index cb12045260e..7ca1ced6b32 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -57,6 +57,8 @@ type options struct {
 	w                 string
 	rebalanceInterval time.Duration
 	tlsClientConfig   *tls.Config
+	auditEnabled      bool
+	auditDir          string
 }
 
 var opts options
@@ -201,6 +203,8 @@ func run() {
 		w:                 Zero.Conf.GetString("wal"),
 		rebalanceInterval: Zero.Conf.GetDuration("rebalance_interval"),
 		tlsClientConfig:   tlsConf,
+		auditEnabled:      Zero.Conf.GetBool("audit_enabled"),
+		auditDir:          Zero.Conf.GetString("audit_dir"),
 	}
 	glog.Infof("Setting Config to: %+v", opts)
 	x.WorkerConfig.Parse(Zero.Conf)
@@ -224,10 +228,9 @@ func run() {
 
 	wd, err := filepath.Abs(opts.w)
 	x.Check(err)
-	ad, err := filepath.Abs(Zero.Conf.GetString("audit_dir"))
+	ad, err := filepath.Abs(opts.auditDir)
 	x.Check(err)
-	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').",
-		Zero.Conf.Get("audit_dir"))
+	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", opts.auditDir)
 
 	if opts.rebalanceInterval <= 0 {
 		log.Fatalf("ERROR: Rebalance interval must be greater than zero. Found: %d",
@@ -279,13 +282,6 @@ func run() {
 	http.HandleFunc("/jemalloc", x.JemallocHandler)
 	zpages.Handle(http.DefaultServeMux, "/z")
 
-	audit.InitAuditorIfNecessary(Zero.Conf, func() bool {
-		if st.zero.state == nil {
-			return false
-		}
-		return st.zero.state.GetLicense().GetEnabled()
-	})
-
 	// This must be here. It does not work if placed before Grpc init.
 	x.Check(st.node.initAndStartNode())
 
@@ -345,5 +341,5 @@ func run() {
 	st.zero.orc.close()
 	glog.Infoln("All done. Goodbye!")
 	audit.Close()
-	glog.Infoln("audit logs are closed")
+	glog.Infoln("audit logs if enabled are closed.")
 }
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index b40a41f67df..41267b24fc1 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -27,33 +27,44 @@ type AuditEvent struct {
 const (
 	UnauthorisedUser = "UnauthorisedUser"
 	UnknownUser      = "UnknownUser"
+	PoorManAuth      = "PoorManAuth"
 	Grpc             = "Grpc"
 	Http             = "Http"
 )
 
-var auditor *auditLogger
+var auditor *auditLogger = &auditLogger{}
 
 type auditLogger struct {
 	log  *x.Logger
 	tick *time.Ticker
 }
 
+// InitAuditorIfNecessary accepts conf and enterprise edition check function.
+// This method keep tracks whether cluster is part of enterprise edition or not.
+// It pools eeEnabled function every five minutes to check if the license is still valid or not.
 func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
 	if !conf.GetBool("audit_enabled") {
 		return
 	}
-	auditor = &auditLogger{
-		tick: time.NewTicker(time.Minute * 5),
-	}
 	if eeEnabled() {
-		auditor.log = initlog(conf.GetString("audit_dir"))
-		atomic.StoreUint32(&auditEnabled, 1)
-		glog.Infoln("audit logs are enabled")
+		InitAuditor(true, conf.GetString("audit_dir"))
 	}
-
+	auditor.tick = time.NewTicker(time.Minute * 5)
 	go trackIfEEValid(eeEnabled, conf.GetString("audit_dir"))
 }
 
+// InitAuditor initializes the auditor.
+// This method doesnt keep track of whether cluster is part of enterprise edition or not.
+// Client has to keep track of that.
+func InitAuditor(enabled bool, dir string) {
+	if !enabled{
+		return
+	}
+	auditor.log = initlog(dir)
+	atomic.StoreUint32(&auditEnabled, 1)
+	glog.Infoln("audit logs are enabled")
+}
+
 func initlog(dir string) *x.Logger {
 	logger, err := x.InitLogger(dir, "dgraph_audit.log")
 	if err != nil {
@@ -86,12 +97,15 @@ func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
 	}
 }
 
+// Close stops the ticker and sync the pending logs in buffer.
+// It also sets the log to nil, because its being called by zero when license expires.
+// If license added, InitLogger will take care of the file.
 func Close() {
-	if auditor == nil {
-		return
+	if auditor.tick != nil {
+		auditor.tick.Stop()
 	}
-	auditor.tick.Stop()
 	auditor.log.Sync()
+	auditor.log = nil
 }
 
 func (a *auditLogger) AuditEvent(event *AuditEvent) {
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 55ef002bcd4..13ec2b63784 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -60,32 +60,32 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 }
 
 func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error) {
-	var code codes.Code
-	if serr, ok := status.FromError(err); !ok {
-		code = codes.Unknown
-	} else {
-		code = serr.Code()
-	}
 	clientHost := ""
 	if p, ok := peer.FromContext(ctx); ok {
 		clientHost = p.Addr.String()
 	}
 
-	token := ""
+	userId := ""
 	if md, ok := metadata.FromIncomingContext(ctx); ok {
 		if t := md.Get("accessJwt"); len(t) > 0 {
-			token = t[0]
+			userId = getUserId(t[0], false)
+		} else if t := md.Get("auth-token"); len(t) > 0 {
+			userId = getUserId(t[0], true)
 		}
 	}
 
+	cd := codes.Unknown
+	if serr, ok := status.FromError(err); ok {
+		cd = serr.Code()
+	}
 	auditor.AuditEvent(&AuditEvent{
-		User:       getUserId(token),
+		User:       userId,
 		ServerHost: x.WorkerConfig.MyAddr,
 		ClientHost: clientHost,
 		Endpoint:   info.FullMethod,
 		ReqType:    Grpc,
 		Req:        fmt.Sprintf("%v", req),
-		Status:     code.String(),
+		Status:     cd.String(),
 	})
 }
 
@@ -94,8 +94,17 @@ func auditHttp(w *ResponseWriter, r *http.Request) {
 	if err != nil {
 		rb = []byte(err.Error())
 	}
+
+	userId := ""
+	if token := r.Header.Get("X-Dgraph-AccessToken"); token != "" {
+		userId = getUserId(token, false)
+	} else if token := r.Header.Get("X-Dgraph-AuthToken"); token != "" {
+		userId = getUserId(token, true)
+	} else {
+		userId = getUserId("", false)
+	}
 	auditor.AuditEvent(&AuditEvent{
-		User:        getUserId(r.Header.Get("X-Dgraph-AccessToken")),
+		User:        userId,
 		ServerHost:  x.WorkerConfig.MyAddr,
 		ClientHost:  r.RemoteAddr,
 		Endpoint:    r.URL.Path,
@@ -106,7 +115,10 @@ func auditHttp(w *ResponseWriter, r *http.Request) {
 	})
 }
 
-func getUserId(token string) string {
+func getUserId(token string, poorman bool) string {
+	if poorman {
+		return PoorManAuth
+	}
 	var userId string
 	var err error
 	if token == "" {

From 9b5c43e56134d219209e8907b737a8736da551f5 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Fri, 15 Jan 2021 20:05:45 +0530
Subject: [PATCH 14/31] fixing zero init audit process

---
 dgraph/cmd/zero/raft.go | 4 ++--
 ee/audit/audit.go       | 7 ++-----
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/dgraph/cmd/zero/raft.go b/dgraph/cmd/zero/raft.go
index 3a8930880f4..b06648e3b96 100644
--- a/dgraph/cmd/zero/raft.go
+++ b/dgraph/cmd/zero/raft.go
@@ -402,8 +402,8 @@ func (n *node) applyProposal(e raftpb.Entry) (uint64, error) {
 		// Check expiry and set enabled accordingly.
 		expiry := time.Unix(state.License.ExpiryTs, 0).UTC()
 		state.License.Enabled = time.Now().UTC().Before(expiry)
-		if state.License.Enabled {
-			audit.InitAuditor(opts.auditEnabled, opts.auditDir)
+		if state.License.Enabled && opts.auditEnabled {
+			audit.InitAuditor(opts.auditDir)
 		}
 	}
 	if p.Snapshot != nil {
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index 41267b24fc1..710756d67bd 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -47,7 +47,7 @@ func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
 		return
 	}
 	if eeEnabled() {
-		InitAuditor(true, conf.GetString("audit_dir"))
+		InitAuditor(conf.GetString("audit_dir"))
 	}
 	auditor.tick = time.NewTicker(time.Minute * 5)
 	go trackIfEEValid(eeEnabled, conf.GetString("audit_dir"))
@@ -56,10 +56,7 @@ func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
 // InitAuditor initializes the auditor.
 // This method doesnt keep track of whether cluster is part of enterprise edition or not.
 // Client has to keep track of that.
-func InitAuditor(enabled bool, dir string) {
-	if !enabled{
-		return
-	}
+func InitAuditor(dir string) {
 	auditor.log = initlog(dir)
 	atomic.StoreUint32(&auditEnabled, 1)
 	glog.Infoln("audit logs are enabled")

From d7adbf8a03e9f11b846f48f19aaaa25bcad07e2f Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Mon, 18 Jan 2021 12:44:11 +0530
Subject: [PATCH 15/31] adding licence file + build tags

---
 dgraph/cmd/alpha/run.go |   2 +-
 ee/audit/audit.go       | 124 +++++++----------------
 ee/audit/audit_ee.go    | 216 ++++++++++++++++++++++++++++++++++++++++
 ee/audit/interceptor.go | 121 ++++++----------------
 worker/config.go        |   9 +-
 x/config.go             |   2 -
 x/jwt_helper.go         |  16 +++
 x/logger.go             |  16 +++
 8 files changed, 319 insertions(+), 187 deletions(-)
 create mode 100644 ee/audit/audit_ee.go

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index caf78e76022..0ab23a970a7 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -613,6 +613,7 @@ func run() {
 		MutationsMode: worker.AllowMutations,
 		AuthToken:     Alpha.Conf.GetString("auth_token"),
 		AuditEnabled:  Alpha.Conf.GetBool("audit_enabled"),
+		AuditDir:      Alpha.Conf.GetString("audit_dir"),
 	}
 
 	secretFile := Alpha.Conf.GetString("acl_secret_file")
@@ -675,7 +676,6 @@ func run() {
 		TLSClientConfig:      tlsClientConf,
 		TLSServerConfig:      tlsServerConf,
 		HmacSecret:           opts.HmacSecret,
-		AuditDir:             Alpha.Conf.GetString("audit_dir"),
 	}
 	x.WorkerConfig.Parse(Alpha.Conf)
 	x.CheckFlag(x.WorkerConfig.Raft, "group", "idx", "learner")
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index 710756d67bd..1ca78dc0b55 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -1,117 +1,61 @@
+// +build oss
+
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package audit
 
 import (
-	"sync/atomic"
-	"time"
+	"context"
+	"google.golang.org/grpc"
+	"net/http"
 
 	"github.com/spf13/viper"
-
-	"github.com/golang/glog"
-
-	"github.com/dgraph-io/dgraph/x"
 )
 
 var auditEnabled uint32
 
 type AuditEvent struct {
-	User        string
-	ServerHost  string
-	ClientHost  string
-	Endpoint    string
-	ReqType     string
-	Req         string
-	Status      string
-	QueryParams map[string][]string
-}
 
-const (
-	UnauthorisedUser = "UnauthorisedUser"
-	UnknownUser      = "UnknownUser"
-	PoorManAuth      = "PoorManAuth"
-	Grpc             = "Grpc"
-	Http             = "Http"
-)
+}
 
-var auditor *auditLogger = &auditLogger{}
+var auditor *auditLogger
 
 type auditLogger struct {
-	log  *x.Logger
-	tick *time.Ticker
 }
 
-// InitAuditorIfNecessary accepts conf and enterprise edition check function.
-// This method keep tracks whether cluster is part of enterprise edition or not.
-// It pools eeEnabled function every five minutes to check if the license is still valid or not.
 func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
-	if !conf.GetBool("audit_enabled") {
-		return
-	}
-	if eeEnabled() {
-		InitAuditor(conf.GetString("audit_dir"))
-	}
-	auditor.tick = time.NewTicker(time.Minute * 5)
-	go trackIfEEValid(eeEnabled, conf.GetString("audit_dir"))
+	return
 }
 
-// InitAuditor initializes the auditor.
-// This method doesnt keep track of whether cluster is part of enterprise edition or not.
-// Client has to keep track of that.
 func InitAuditor(dir string) {
-	auditor.log = initlog(dir)
-	atomic.StoreUint32(&auditEnabled, 1)
-	glog.Infoln("audit logs are enabled")
+	return
 }
-
-func initlog(dir string) *x.Logger {
-	logger, err := x.InitLogger(dir, "dgraph_audit.log")
-	if err != nil {
-		glog.Errorf("error while initiating auditor %v", err)
-		return nil
-	}
-	return logger
+func Close() {
+	return
 }
 
-// trackIfEEValid tracks enterprise license of the cluster.
-// Right now alpha doesn't know about the enterprise/licence.
-// That's why we needed to track if the current node is part of enterprise edition cluster
-func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
-	for {
-		select {
-		case <-auditor.tick.C:
-			if !eeEnabledFunc() && atomic.CompareAndSwapUint32(&auditEnabled, 1, 0) {
-				glog.Infof("audit logs are disabled")
-				auditor.log.Sync()
-				auditor.log = nil
-				continue
-			}
-
-			if atomic.LoadUint32(&auditEnabled) != 1 {
-				auditor.log = initlog(dir)
-				atomic.StoreUint32(&auditEnabled, 1)
-				glog.Infof("audit logs are enabled")
-			}
-		}
-	}
+func (a *auditLogger) Audit(event *AuditEvent) {
+	return
 }
 
-// Close stops the ticker and sync the pending logs in buffer.
-// It also sets the log to nil, because its being called by zero when license expires.
-// If license added, InitLogger will take care of the file.
-func Close() {
-	if auditor.tick != nil {
-		auditor.tick.Stop()
-	}
-	auditor.log.Sync()
-	auditor.log = nil
+func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error) {
+	return
 }
 
-func (a *auditLogger) AuditEvent(event *AuditEvent) {
-	a.log.AuditI(event.Endpoint,
-		"user", event.User,
-		"server", event.ServerHost,
-		"client", event.ClientHost,
-		"req_type", event.ReqType,
-		"req_body", event.Req,
-		"query_param", event.QueryParams,
-		"status", event.Status)
-}
+func auditHttp(w *ResponseWriter, r *http.Request) {
+	return
+}
\ No newline at end of file
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
new file mode 100644
index 00000000000..5284c52db6b
--- /dev/null
+++ b/ee/audit/audit_ee.go
@@ -0,0 +1,216 @@
+// +build !oss
+
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package audit
+
+import (
+	"context"
+	"fmt"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
+	"io/ioutil"
+	"net/http"
+	"sync/atomic"
+	"time"
+
+	"github.com/dgraph-io/dgraph/x"
+	"github.com/golang/glog"
+	"github.com/spf13/viper"
+)
+
+var auditEnabled uint32
+
+type AuditEvent struct {
+	User        string
+	ServerHost  string
+	ClientHost  string
+	Endpoint    string
+	ReqType     string
+	Req         string
+	Status      string
+	QueryParams map[string][]string
+}
+
+const (
+	UnauthorisedUser = "UnauthorisedUser"
+	UnknownUser      = "UnknownUser"
+	PoorManAuth      = "PoorManAuth"
+	Grpc             = "Grpc"
+	Http             = "Http"
+)
+
+var auditor *auditLogger = &auditLogger{}
+
+type auditLogger struct {
+	log  *x.Logger
+	tick *time.Ticker
+}
+
+// InitAuditorIfNecessary accepts conf and enterprise edition check function.
+// This method keep tracks whether cluster is part of enterprise edition or not.
+// It pools eeEnabled function every five minutes to check if the license is still valid or not.
+func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
+	if !conf.GetBool("audit_enabled") {
+		return
+	}
+	if eeEnabled() {
+		InitAuditor(conf.GetString("audit_dir"))
+	}
+	auditor.tick = time.NewTicker(time.Minute * 5)
+	go trackIfEEValid(eeEnabled, conf.GetString("audit_dir"))
+}
+
+// InitAuditor initializes the auditor.
+// This method doesnt keep track of whether cluster is part of enterprise edition or not.
+// Client has to keep track of that.
+func InitAuditor(dir string) {
+	auditor.log = initlog(dir)
+	atomic.StoreUint32(&auditEnabled, 1)
+	glog.Infoln("audit logs are enabled")
+}
+
+func initlog(dir string) *x.Logger {
+	logger, err := x.InitLogger(dir, "dgraph_audit.log")
+	if err != nil {
+		glog.Errorf("error while initiating auditor %v", err)
+		return nil
+	}
+	return logger
+}
+
+// trackIfEEValid tracks enterprise license of the cluster.
+// Right now alpha doesn't know about the enterprise/licence.
+// That's why we needed to track if the current node is part of enterprise edition cluster
+func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
+	for {
+		select {
+		case <-auditor.tick.C:
+			if !eeEnabledFunc() && atomic.CompareAndSwapUint32(&auditEnabled, 1, 0) {
+				glog.Infof("audit logs are disabled")
+				auditor.log.Sync()
+				auditor.log = nil
+				continue
+			}
+
+			if atomic.LoadUint32(&auditEnabled) != 1 {
+				auditor.log = initlog(dir)
+				atomic.StoreUint32(&auditEnabled, 1)
+				glog.Infof("audit logs are enabled")
+			}
+		}
+	}
+}
+
+// Close stops the ticker and sync the pending logs in buffer.
+// It also sets the log to nil, because its being called by zero when license expires.
+// If license added, InitLogger will take care of the file.
+func Close() {
+	if auditor.tick != nil {
+		auditor.tick.Stop()
+	}
+	auditor.log.Sync()
+	auditor.log = nil
+}
+
+func (a *auditLogger) Audit(event *AuditEvent) {
+	a.log.AuditI(event.Endpoint,
+		"user", event.User,
+		"server", event.ServerHost,
+		"client", event.ClientHost,
+		"req_type", event.ReqType,
+		"req_body", event.Req,
+		"query_param", event.QueryParams,
+		"status", event.Status)
+}
+
+func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error) {
+	clientHost := ""
+	if p, ok := peer.FromContext(ctx); ok {
+		clientHost = p.Addr.String()
+	}
+
+	userId := ""
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		if t := md.Get("accessJwt"); len(t) > 0 {
+			userId = getUserId(t[0], false)
+		} else if t := md.Get("auth-token"); len(t) > 0 {
+			userId = getUserId(t[0], true)
+		}
+	}
+
+	cd := codes.Unknown
+	if serr, ok := status.FromError(err); ok {
+		cd = serr.Code()
+	}
+	auditor.Audit(&AuditEvent{
+		User:       userId,
+		ServerHost: x.WorkerConfig.MyAddr,
+		ClientHost: clientHost,
+		Endpoint:   info.FullMethod,
+		ReqType:    Grpc,
+		Req:        fmt.Sprintf("%v", req),
+		Status:     cd.String(),
+	})
+}
+
+func auditHttp(w *ResponseWriter, r *http.Request) {
+	rb, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		rb = []byte(err.Error())
+	}
+
+	userId := ""
+	if token := r.Header.Get("X-Dgraph-AccessToken"); token != "" {
+		userId = getUserId(token, false)
+	} else if token := r.Header.Get("X-Dgraph-AuthToken"); token != "" {
+		userId = getUserId(token, true)
+	} else {
+		userId = getUserId("", false)
+	}
+	auditor.Audit(&AuditEvent{
+		User:        userId,
+		ServerHost:  x.WorkerConfig.MyAddr,
+		ClientHost:  r.RemoteAddr,
+		Endpoint:    r.URL.Path,
+		ReqType:     Http,
+		Req:         string(rb),
+		Status:      http.StatusText(w.statusCode),
+		QueryParams: r.URL.Query(),
+	})
+}
+
+func getUserId(token string, poorman bool) string {
+	if poorman {
+		return PoorManAuth
+	}
+	var userId string
+	var err error
+	if token == "" {
+		if x.WorkerConfig.AclEnabled {
+			userId = UnauthorisedUser
+		}
+	} else {
+		if userId, err = x.ExtractUserName(token); err != nil {
+			userId = UnknownUser
+		}
+	}
+	return userId
+}
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 13ec2b63784..e94b3f7828a 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -1,37 +1,52 @@
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package audit
 
 import (
 	"bytes"
 	"context"
-	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
 	"strings"
 	"sync/atomic"
 
-	"github.com/dgraph-io/dgraph/x"
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/metadata"
-	"google.golang.org/grpc/peer"
-	"google.golang.org/grpc/status"
 )
 
 func AuditRequestGRPC(ctx context.Context, req interface{},
 	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 	skip := func(method string) bool {
-		skipApis := []string{"Heartbeat", "RaftMessage", "JoinCluster", "IsPeer", // raft server
-			"StreamMembership", "UpdateMembership", "Oracle", // zero server
-			"Check", "Watch", // health server
-		}
-
-		for _, api := range skipApis {
-			if strings.HasSuffix(method, api) {
-				return true
-			}
+		skipApis := map[string]bool{
+			// raft server
+			"Heartbeat":   true,
+			"RaftMessage": true,
+			"JoinCluster": true,
+			"IsPeer":      true,
+			// zero server
+			"StreamMembership": true,
+			"UpdateMembership": true,
+			"Oracle":           true,
+			"Timestamps":       true,
+			// health server
+			"Check": true,
+			"Watch": true,
 		}
-		return false
+		return skipApis[info.FullMethod[strings.LastIndex(info.FullMethod, "/")+1:]]
 	}
 
 	if atomic.LoadUint32(&auditEnabled) == 0 || skip(info.FullMethod) {
@@ -59,80 +74,6 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 	})
 }
 
-func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error) {
-	clientHost := ""
-	if p, ok := peer.FromContext(ctx); ok {
-		clientHost = p.Addr.String()
-	}
-
-	userId := ""
-	if md, ok := metadata.FromIncomingContext(ctx); ok {
-		if t := md.Get("accessJwt"); len(t) > 0 {
-			userId = getUserId(t[0], false)
-		} else if t := md.Get("auth-token"); len(t) > 0 {
-			userId = getUserId(t[0], true)
-		}
-	}
-
-	cd := codes.Unknown
-	if serr, ok := status.FromError(err); ok {
-		cd = serr.Code()
-	}
-	auditor.AuditEvent(&AuditEvent{
-		User:       userId,
-		ServerHost: x.WorkerConfig.MyAddr,
-		ClientHost: clientHost,
-		Endpoint:   info.FullMethod,
-		ReqType:    Grpc,
-		Req:        fmt.Sprintf("%v", req),
-		Status:     cd.String(),
-	})
-}
-
-func auditHttp(w *ResponseWriter, r *http.Request) {
-	rb, err := ioutil.ReadAll(r.Body)
-	if err != nil {
-		rb = []byte(err.Error())
-	}
-
-	userId := ""
-	if token := r.Header.Get("X-Dgraph-AccessToken"); token != "" {
-		userId = getUserId(token, false)
-	} else if token := r.Header.Get("X-Dgraph-AuthToken"); token != "" {
-		userId = getUserId(token, true)
-	} else {
-		userId = getUserId("", false)
-	}
-	auditor.AuditEvent(&AuditEvent{
-		User:        userId,
-		ServerHost:  x.WorkerConfig.MyAddr,
-		ClientHost:  r.RemoteAddr,
-		Endpoint:    r.URL.Path,
-		ReqType:     Http,
-		Req:         string(rb),
-		Status:      http.StatusText(w.statusCode),
-		QueryParams: r.URL.Query(),
-	})
-}
-
-func getUserId(token string, poorman bool) string {
-	if poorman {
-		return PoorManAuth
-	}
-	var userId string
-	var err error
-	if token == "" {
-		if x.WorkerConfig.AclEnabled {
-			userId = UnauthorisedUser
-		}
-	} else {
-		if userId, err = x.ExtractUserName(token); err != nil {
-			userId = UnknownUser
-		}
-	}
-	return userId
-}
-
 type ResponseWriter struct {
 	http.ResponseWriter
 	statusCode int
diff --git a/worker/config.go b/worker/config.go
index d9256459c31..490aca142ed 100644
--- a/worker/config.go
+++ b/worker/config.go
@@ -72,6 +72,7 @@ type Options struct {
 	CacheMb int64
 
 	AuditEnabled bool
+	AuditDir string
 }
 
 // Config holds an instance of the server options..
@@ -96,12 +97,12 @@ func (opt *Options) validate() {
 	x.Check(err)
 	td, err := filepath.Abs(x.WorkerConfig.TmpDir)
 	x.Check(err)
-	ad, err := filepath.Abs(x.WorkerConfig.AuditDir)
+	ad, err := filepath.Abs(opt.AuditDir)
 	x.Check(err)
 	x.AssertTruef(pd != wd, "Posting and WAL directory cannot be the same ('%s').", opt.PostingDir)
 	x.AssertTruef(pd != td, "Posting and Tmp directory cannot be the same ('%s').", opt.PostingDir)
 	x.AssertTruef(wd != td, "WAL and Tmp directory cannot be the same ('%s').", opt.WALDir)
-	x.AssertTruef(ad != pd, "Posting and Audit Directory cannot be the same ('%s').", x.WorkerConfig.AuditDir)
-	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", x.WorkerConfig.AuditDir)
-	x.AssertTruef(ad != td, "Tmp and Audit directory cannot be the same ('%s').", x.WorkerConfig.AuditDir)
+	x.AssertTruef(ad != pd, "Posting and Audit Directory cannot be the same ('%s').", opt.AuditDir)
+	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", opt.AuditDir)
+	x.AssertTruef(ad != td, "Tmp and Audit directory cannot be the same ('%s').", opt.AuditDir)
 }
diff --git a/x/config.go b/x/config.go
index 2808c57a135..607848ba078 100644
--- a/x/config.go
+++ b/x/config.go
@@ -112,8 +112,6 @@ type WorkerOptions struct {
 	LogRequest int32
 	// If true, we should call msync or fsync after every write to survive hard reboots.
 	HardSync bool
-
-	AuditDir string
 }
 
 // WorkerConfig stores the global instance of the worker package's options.
diff --git a/x/jwt_helper.go b/x/jwt_helper.go
index 34efa2f9b9e..a8b6d3472aa 100644
--- a/x/jwt_helper.go
+++ b/x/jwt_helper.go
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package x
 
 import (
diff --git a/x/logger.go b/x/logger.go
index faaea4c3fb5..55f60be5b74 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package x
 
 import (

From 0710595b3c7de1c9261a16386ec4ddb2581940fb Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Tue, 19 Jan 2021 10:51:52 +0530
Subject: [PATCH 16/31] adding dgraph audit tool and logs encryption

---
 dgraph/cmd/root_ee.go                   |   2 +
 dgraph/cmd/zero/raft.go                 |   2 +-
 dgraph/cmd/zero/run.go                  |   5 +
 ee/audit/audit_ee.go                    |  40 +++----
 ee/audit/encryption_test.go             | 153 ++++++++++++++++++++++++
 ee/audit/interceptor.go                 |   1 +
 ee/audit/run.go                         |  33 +++++
 ee/audit/run_ee.go                      | 116 ++++++++++++++++++
 ee/audit/test_encryption_chunks_cbc.txt | Bin 0 -> 4800 bytes
 ee/audit/test_out.txt                   | Bin 0 -> 4768 bytes
 go.mod                                  |   1 +
 x/log_writer.go                         | 120 +++++++++++++++++++
 x/logger.go                             |  10 +-
 13 files changed, 459 insertions(+), 24 deletions(-)
 create mode 100644 ee/audit/encryption_test.go
 create mode 100644 ee/audit/run.go
 create mode 100644 ee/audit/run_ee.go
 create mode 100644 ee/audit/test_encryption_chunks_cbc.txt
 create mode 100644 ee/audit/test_out.txt
 create mode 100644 x/log_writer.go

diff --git a/dgraph/cmd/root_ee.go b/dgraph/cmd/root_ee.go
index b1bdc727ed2..5660d8c3450 100644
--- a/dgraph/cmd/root_ee.go
+++ b/dgraph/cmd/root_ee.go
@@ -14,6 +14,7 @@ package cmd
 
 import (
 	acl "github.com/dgraph-io/dgraph/ee/acl"
+	"github.com/dgraph-io/dgraph/ee/audit"
 	"github.com/dgraph-io/dgraph/ee/backup"
 )
 
@@ -24,5 +25,6 @@ func init() {
 		&backup.LsBackup,
 		&backup.ExportBackup,
 		&acl.CmdAcl,
+		&audit.CmdAudit,
 	)
 }
diff --git a/dgraph/cmd/zero/raft.go b/dgraph/cmd/zero/raft.go
index b06648e3b96..94749e92384 100644
--- a/dgraph/cmd/zero/raft.go
+++ b/dgraph/cmd/zero/raft.go
@@ -403,7 +403,7 @@ func (n *node) applyProposal(e raftpb.Entry) (uint64, error) {
 		expiry := time.Unix(state.License.ExpiryTs, 0).UTC()
 		state.License.Enabled = time.Now().UTC().Before(expiry)
 		if state.License.Enabled && opts.auditEnabled {
-			audit.InitAuditor(opts.auditDir)
+			audit.InitAuditor(opts.auditDir, opts.encryptionKey)
 		}
 	}
 	if p.Snapshot != nil {
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index 7ca1ced6b32..d2910e57078 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -59,6 +59,7 @@ type options struct {
 	tlsClientConfig   *tls.Config
 	auditEnabled      bool
 	auditDir          string
+	encryptionKey     []byte
 }
 
 var opts options
@@ -104,6 +105,7 @@ instances to achieve high-availability.
 	flag.Bool("audit_enabled", false, "Set to true to enable audit logs.")
 	flag.String("audit_dir", "za", "Set path to directory where to save the audit logs.")
 
+	enc.RegisterFlags(flag)
 	// TLS configurations
 	x.RegisterServerTLSFlags(flag)
 }
@@ -194,6 +196,8 @@ func run() {
 	x.PrintVersion()
 	tlsConf, err := x.LoadClientTLSConfigForInternalPort(Zero.Conf)
 	x.Check(err)
+
+	key, err := enc.ReadKey(Zero.Conf)
 	opts = options{
 		bindall:           Zero.Conf.GetBool("bindall"),
 		portOffset:        Zero.Conf.GetInt("port_offset"),
@@ -205,6 +209,7 @@ func run() {
 		tlsClientConfig:   tlsConf,
 		auditEnabled:      Zero.Conf.GetBool("audit_enabled"),
 		auditDir:          Zero.Conf.GetString("audit_dir"),
+		encryptionKey:     key,
 	}
 	glog.Infof("Setting Config to: %+v", opts)
 	x.WorkerConfig.Parse(Zero.Conf)
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index 5284c52db6b..5d59a45436d 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -1,19 +1,13 @@
 // +build !oss
 
 /*
- * Copyright 2017-2021 Dgraph Labs, Inc. and Contributors
+ * Copyright 2021 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.
- * You may obtain a copy of the License at
+ * Licensed under the Dgraph Community License (the "License"); you
+ * may not use this file except in compliance with the License. You
+ * may obtain a copy of the License at
  *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ *     https://github.com/dgraph-io/dgraph/blob/master/licenses/DCL.txt
  */
 
 package audit
@@ -21,6 +15,7 @@ package audit
 import (
 	"context"
 	"fmt"
+	"github.com/dgraph-io/dgraph/ee/enc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -71,24 +66,29 @@ func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
 	if !conf.GetBool("audit_enabled") {
 		return
 	}
+	key, err := enc.ReadKey(conf)
+	if err != nil {
+		return
+	}
 	if eeEnabled() {
-		InitAuditor(conf.GetString("audit_dir"))
+
+		InitAuditor(conf.GetString("audit_dir"), key)
 	}
 	auditor.tick = time.NewTicker(time.Minute * 5)
-	go trackIfEEValid(eeEnabled, conf.GetString("audit_dir"))
+	go trackIfEEValid(eeEnabled, conf.GetString("audit_dir"), key)
 }
 
 // InitAuditor initializes the auditor.
 // This method doesnt keep track of whether cluster is part of enterprise edition or not.
 // Client has to keep track of that.
-func InitAuditor(dir string) {
-	auditor.log = initlog(dir)
+func InitAuditor(dir string, key []byte) {
+	auditor.log = initlog(dir, key)
 	atomic.StoreUint32(&auditEnabled, 1)
 	glog.Infoln("audit logs are enabled")
 }
 
-func initlog(dir string) *x.Logger {
-	logger, err := x.InitLogger(dir, "dgraph_audit.log")
+func initlog(dir string, key []byte) *x.Logger {
+	logger, err := x.InitLogger(dir, "dgraph_audit.log", key)
 	if err != nil {
 		glog.Errorf("error while initiating auditor %v", err)
 		return nil
@@ -99,7 +99,7 @@ func initlog(dir string) *x.Logger {
 // trackIfEEValid tracks enterprise license of the cluster.
 // Right now alpha doesn't know about the enterprise/licence.
 // That's why we needed to track if the current node is part of enterprise edition cluster
-func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
+func trackIfEEValid(eeEnabledFunc func() bool, dir string, key []byte) {
 	for {
 		select {
 		case <-auditor.tick.C:
@@ -111,7 +111,7 @@ func trackIfEEValid(eeEnabledFunc func() bool, dir string) {
 			}
 
 			if atomic.LoadUint32(&auditEnabled) != 1 {
-				auditor.log = initlog(dir)
+				auditor.log = initlog(dir, key)
 				atomic.StoreUint32(&auditEnabled, 1)
 				glog.Infof("audit logs are enabled")
 			}
@@ -166,7 +166,7 @@ func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
 		ClientHost: clientHost,
 		Endpoint:   info.FullMethod,
 		ReqType:    Grpc,
-		Req:        fmt.Sprintf("%v", req),
+		Req:        fmt.Sprintf("%+v", req),
 		Status:     cd.String(),
 	})
 }
diff --git a/ee/audit/encryption_test.go b/ee/audit/encryption_test.go
new file mode 100644
index 00000000000..e5b529b8596
--- /dev/null
+++ b/ee/audit/encryption_test.go
@@ -0,0 +1,153 @@
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package audit
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"fmt"
+	"github.com/stretchr/testify/require"
+	"io"
+	"io/ioutil"
+	"os"
+	"testing"
+)
+
+var key = []byte("12345678901234561234567890123456")
+var entry = []byte("abcdefghijklmnopqrstuvwxyz")
+
+func TestEncryptionCBC(t *testing.T) {
+	file, err := os.OpenFile("test_encryption_chunks_cbc.txt", os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
+		0600)
+	require.Nil(t, err)
+	salt := make([]byte, 8) // Generate an 8 byte salt
+	rand.Read(salt)
+	data := make([]byte, len(entry) + aes.BlockSize)
+	copy(data[0:], "Salted__")
+	copy(data[8:], salt)
+
+	for i := 0; i<100; i ++ {
+		copy(data[aes.BlockSize:], entry)
+		padded, err := pkcs7Pad(data, aes.BlockSize)
+		c, err := aes.NewCipher(key)
+		require.Nil(t, err)
+		iv := make([]byte, aes.BlockSize)
+		rand.Read(iv)
+		cbc := cipher.NewCBCEncrypter(c, iv)
+		cbc.CryptBlocks(padded[aes.BlockSize:], padded[aes.BlockSize:])
+		file.Write(padded)
+	}
+
+	file.Close()
+}
+
+// pkcs7Pad appends padding.
+func pkcs7Pad(data []byte, blocklen int) ([]byte, error) {
+	if blocklen <= 0 {
+		return nil, fmt.Errorf("invalid blocklen %d", blocklen)
+	}
+	padlen := 1
+	for ((len(data) + padlen) % blocklen) != 0 {
+		padlen++
+	}
+
+	pad := bytes.Repeat([]byte{byte(padlen)}, padlen)
+	return append(data, pad...), nil
+}
+
+func TestEncryptionInChunks(t *testing.T) {
+	file, err := os.OpenFile("test_encryption_chunks.txt", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	require.Nil(t, err)
+	iv := make([]byte, aes.BlockSize)
+	rand.Read(iv)
+
+	block, err := aes.NewCipher(key)
+	require.Nil(t, err)
+
+	stream := cipher.NewCTR(block, iv)
+	for i:=0; i< 100; i++ {
+		e := []byte("abcdefghijklmnopqrstuvwxyz")
+		stream.XORKeyStream(e, e)
+		file.Write(e)
+	}
+	file.Write(iv)
+}
+
+func TestDecryptionInChunks(t *testing.T) {
+	cipherText, err := os.Open("test_encryption_chunks.txt")
+	require.Nil(t, err)
+	iv := make([]byte, aes.BlockSize)
+	stat, _ := cipherText.Stat()
+	msg := stat.Size() - int64(len(iv))
+	_, err = cipherText.ReadAt(iv, msg)
+	require.Nil(t, err)
+	outfile, err := os.OpenFile("test_encryption_chunks_out.txt",
+		os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
+		0600)
+	require.Nil(t, err)
+
+	block, err := aes.NewCipher(key)
+	require.Nil(t, err)
+
+	buf := make([]byte, 26)
+	stream := cipher.NewCTR(block, iv)
+
+	for {
+		n, err := cipherText.Read(buf)
+		if n > 0 {
+			if msg == 0 {
+				break
+			}
+			msg = msg - int64(n)
+			stream.XORKeyStream(buf, buf)
+			outfile.Write(buf)
+		}
+
+		if err == io.EOF { break }
+		require.Nil(t, err)
+	}
+}
+
+func TestEncryptionAtOnce(t *testing.T) {
+	file, err := os.OpenFile("test_encryption_once.txt", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	require.Nil(t, err)
+	block, err := aes.NewCipher(key)
+	require.Nil(t, err)
+	iv := make([]byte, aes.BlockSize)
+	rand.Read(iv)
+	gcmCipher := cipher.NewCBCEncrypter(block, iv)
+	repeat := bytes.Repeat(entry, 100)
+
+	gcmCipher.CryptBlocks(repeat, repeat)
+	file.Write(repeat)
+}
+
+func TestDecryptionAtOnce(t *testing.T) {
+	cipherText, err := ioutil.ReadFile("test_encryption_once.txt")
+	require.Nil(t, err)
+	block, err := aes.NewCipher(key)
+	require.Nil(t, err)
+	gcm, err := cipher.NewGCM(block)
+	require.Nil(t, err)
+	nonce := cipherText[:gcm.NonceSize()]
+	cipherText = cipherText[gcm.NonceSize():]
+	plaintext, err := gcm.Open(nil, nonce, cipherText, nil)
+	require.Nil(t, err)
+	fmt.Println(string(plaintext))
+}
\ No newline at end of file
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index e94b3f7828a..60f2563cc7a 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -42,6 +42,7 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 			"UpdateMembership": true,
 			"Oracle":           true,
 			"Timestamps":       true,
+			"ShouldServe":      true,
 			// health server
 			"Check": true,
 			"Watch": true,
diff --git a/ee/audit/run.go b/ee/audit/run.go
new file mode 100644
index 00000000000..3dc44b7ac39
--- /dev/null
+++ b/ee/audit/run.go
@@ -0,0 +1,33 @@
+// +build oss
+
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package audit
+
+import (
+	"github.com/dgraph-io/dgraph/x"
+	"github.com/spf13/cobra"
+)
+
+var CmdAudit x.SubCommand
+
+func init() {
+	CmdAudit.Cmd = &cobra.Command{
+		Use:   "audit",
+		Short: "Enterprise feature. Not supported in oss version",
+	}
+}
diff --git a/ee/audit/run_ee.go b/ee/audit/run_ee.go
new file mode 100644
index 00000000000..1a4fab90fc6
--- /dev/null
+++ b/ee/audit/run_ee.go
@@ -0,0 +1,116 @@
+// +build !oss
+
+/*
+ * Copyright 2121 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
+ * may obtain a copy of the License at
+ *
+ *     https://github.com/dgraph-io/dgraph/blob/master/licenses/DCL.txt
+ */
+
+package audit
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"github.com/dgraph-io/dgraph/ee/enc"
+	"github.com/dgraph-io/dgraph/x"
+	"github.com/golang/glog"
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+	"os"
+)
+
+var CmdAudit x.SubCommand
+
+func init() {
+	CmdAudit.Cmd = &cobra.Command{
+		Use:   "audit",
+		Short: "Dgraph audit tool",
+	}
+
+	subcommands := initSubcommands()
+	for _, sc := range subcommands {
+		CmdAudit.Cmd.AddCommand(sc.Cmd)
+		sc.Conf = viper.New()
+		if err := sc.Conf.BindPFlags(sc.Cmd.Flags()); err != nil {
+			glog.Fatalf("Unable to bind flags for command %v: %v", sc, err)
+		}
+		if err := sc.Conf.BindPFlags(CmdAudit.Cmd.PersistentFlags()); err != nil {
+			glog.Fatalf("Unable to bind persistent flags from audit for command %v: %v", sc, err)
+		}
+		sc.Conf.SetEnvPrefix(sc.EnvPrefix)
+	}
+}
+
+var decryptCmd x.SubCommand
+
+func initSubcommands() []*x.SubCommand {
+	decryptCmd.Cmd = &cobra.Command{
+		Use:   "decrypt",
+		Short: "Run Dgraph Audit tool to decrypt audit files",
+		Run: func(cmd *cobra.Command, args []string) {
+			if err := run(); err != nil {
+				fmt.Printf("%v\n", err)
+				os.Exit(1)
+			}
+		},
+	}
+
+	decFlags := decryptCmd.Cmd.Flags()
+	decFlags.String("in", "", "input file that needs to decrypted.")
+	decFlags.String("out", "audit_log_out.log",
+		"output file to which decrypted output will be dumped.")
+	enc.RegisterFlags(decFlags)
+	return []*x.SubCommand{&decryptCmd}
+}
+
+func run() error {
+	key, err := enc.ReadKey(decryptCmd.Conf)
+	x.Check(err)
+	if key == nil {
+		return errors.New("no encryption key provided")
+	}
+	file, err := os.Open(decryptCmd.Conf.GetString("in"))
+	x.Check(err)
+	defer file.Close()
+
+	outfile, err := os.OpenFile(decryptCmd.Conf.GetString("out"), os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
+		0600)
+	x.Check(err)
+	defer outfile.Close()
+
+	block, err := aes.NewCipher(key)
+	stat, err := os.Stat(decryptCmd.Conf.GetString("in"))
+	x.Check(err)
+
+	var iterator int64 = 0
+	for {
+		if iterator >= stat.Size() {
+			break
+		}
+		l := make([]byte, 4)
+		_, err := file.ReadAt(l, iterator)
+		x.Check(err)
+		iterator = iterator + 4
+		content := make([]byte, binary.BigEndian.Uint32(l))
+		_, err = file.ReadAt(content, iterator)
+		x.Check(err)
+
+		iterator = iterator + int64(binary.BigEndian.Uint32(l))
+		iv := make([]byte, aes.BlockSize)
+		_, err = file.ReadAt(iv, iterator)
+		iterator = iterator + aes.BlockSize
+
+		stream := cipher.NewCTR(block, iv)
+		stream.XORKeyStream(content, content)
+		_, err = outfile.Write(content)
+		x.Check(err)
+	}
+	return nil
+}
diff --git a/ee/audit/test_encryption_chunks_cbc.txt b/ee/audit/test_encryption_chunks_cbc.txt
new file mode 100644
index 0000000000000000000000000000000000000000..1330a96edd168ec59b68f576cd100845c2a5f13f
GIT binary patch
literal 4800
zcmZYBcRUpSAHeZCLe`l@GNR)$Dw%OM8Re{$QQ@q}o=IO6-5tr^d;VPZI(wc&=4G6c
z5RPQaip1~#kN=;K*W>efd>)_A`(t9`;$v@TWd(14a+20nAedI(>U^VUyF=8rDa}pu
z;tooZ6zw9QQ$H`BeS>AyM8yQ~UUPI?zlBa;Yrejor22HG?3z@n>FlRBj7VutbGRzf
z8uauqFQnw+hF@Wy;!jx=3!NcZ`(uku3`#>hbKwLF836sD_rRrFqxNB_jaY%udBY*&
zlCMWyu#zsqdBK8C$?^iwzaM5E-M#4uaWoP>*1Xi5Bp<Btji$shV_07!(aO|P4D@2p
zT%>aC|JYMnr5mL_s7sah1#f8{sn_k?q57Hx-<|}0L`rz&G2`-nsr;>J_K%9!AH04}
zk3Aq1v&4Ski15En27Lfu&2sxr$q0wmVfw?-xfFy8{P^wvXf~WK!^ZmgRO~^o?(^TL
z^A0;hGJGq42Wo60iX+>DOz5VJAg`*}Q*%Y0f_^5=+2;>q+UY7!u)yL<jPvVtx1U!S
z)~8$W^<J`UFEOA;_6FMa){bOrWZB~C?CK5$HP5+Z3sb`{rQuHDK0KPBe?2w5>?wV|
zM6Q|giF4ddOg>}vC7Y?ngr=<gB5ao>5A-;#AM=T|93EK%oox}w%9pLGuKpjTykGdh
zodP07znodW#GtvVazL{d^}-P|ejhc<m^t$^L3hRUyIime!lM2=*xw*kzpo{dX&swI
zi(Zymo?D2G?#DF~T{t(K^a%IeWI#U=v=Me|H#o%O=O~rnfI{wse7><{Ose`rb|ol*
zY2nO%H>nvO?z%o)p>w3iFh$=SoTt9KM2ULi)-_Kiv|zt5g8ksq3R{!}-&9~Q&Xz`+
z!08+p*(e&3-h>o)(0;c&U;uiv*&iJb_D|;DAm&Gsn2WbGg@n^>y7CM$d)2-Br87OC
z-;qEFH{D}*WuJ7@JKj|l=PxSX{}$R2nz&_4gSXY(0DU~q2g-*&+K&As7FtZ%&da&6
znY4KUMZ2J0j@E_Ns0Pr(U#+hwgx&x3h?`h%96wJo=w55~*s0B1Z*e!RC2q8Xo?Umr
zk>XI!^VVgz0$-@~U6-ppI0#4XjP=8G-rAcafj&P662BoLt6JvT@CG%K%tOt+X%SgP
z3S{pm6g1z{_60p_+q}l69~Y8+{sqdprtF7wi7-ERjXw4oj(=3RIXw~dW+5&_+PcF4
z-{)-d@!z$IsD`3splHGHC-~(0j|;RA&};qCv0g*E6DC&E{Z(m*Xptpv+BQSf==%%W
z9g~)NmY`23=a}$wBC!dBm6z>Tm*Fy9#$o{$86HhZm?*3kLn-KISex*PN0b$VxAG3{
zZ3t$>#<!@8liazwOSp#5tpqR7quI$}B%*VYRh~-Kq(HoulZ(;A!Jh+f5@t>QgShI4
zfnFx#ci=mreUtsB^2W`_s;;HHE8WUh+<eVOTzLzBsBD5hhi0sWo>lsi^HLqWEmd9N
z3f=}e>jC8@DMco=mqV>U{}Pr*LzN)E8fz9G*GKH8G=!4GNC+E)iJpb>pkF5vpvQb3
zG_thV(75pTl3kbF0|%$Mz3*uQe{Z5|AV2?FgcXB+Aj{qc&rtKt%svIf!g$a}xXF7o
z%q{he=`f6Vr^~Gu^ps?*qB?VO(2FkLAgR{&<58LmJ<?5<wMGWZ+^q;PVbH7cIRDkp
zdC_UW+k_DqXIfnf<jh%;F#FLGJDYWhc0K^~n)LQ^^n;Hy-drtwZFl)RTfd)s))U%q
z*goUbrOXdc-+`W?Dwi*Fi?D6|nzGG5S+}fMef7v*ZE-jvw)=m_h>#V~k0!;6r#x1*
z7tS8H3cM9&UtgO%&(VpNvlx!8(dC@#1U*^5SP6?!owf7WOX~H0hF0I3MtxfEw^P@4
zbZO@hK6AesSS0%e+`SYRU*6~(GaN(gL=(CFQD0OtSg(9EFM?Ts{kvYIcoSIVf~{Zu
z@Hehg<V}D1?tyiUn4VGUTm_?pbkN^JSZpJwS5wAskLLwMn{{1umxj7x)5HUd_xf3i
zpN)b(@k(QNdiK!7jW<UWBMA!Ur=l$lbr{M^nG~tQ6}Ivg^!l{ay=)gVA!^W*KFa2&
znh1}Lo7}=&`^i-v-tjnQ9OxZ&<9s!Te&Z39ntdk%fBA4*k_LX{i^mtD6s>ejx?Mp3
z-Zjbn&@nFc5vc|9ee7$g=O=H@n;86fj=H>v4Wv~B^yLjH!$&D65BcRgS)Wk$q!LWd
z!OsOehxX&<?Be@x+yXr-g@&q)EE(HWmpvVPSi|u=XA62D2Dvnxd0X%JI42PFdD@F*
z!9{xGZ$mF%IFYv64k2DZbQ>#iXEi-#GU!2gfPPLyypWgTw<ZPO4ScRPwzQg4+psX_
zVJH>)ZT~CUkr(uoYT}8LboOh~R^&LL5qh^s`Q4!0QFI0-s=|@Q<3a}LF*(O-qU<u2
zr!Gu-)qKlDp4f9NE{68-BQF`z18D4-`S>Ka?Z(t5IegPTq?N`0hUkRhBL<D+_E_27
z;EbD=rbb|2M{m3Y^Cz#q6c@xrRG4U0)*TOWZ)!M{wbO{(EQ<~Xy;x>ab=)0vd|U#|
zyFRdc{SwxFkB9#;dG>(UNGRfT81zw|SUj6+(oB??F0}>HaL!RUmXkH-Ix*~XHWVhd
zg9m-EF)AOkpp3%@h;E4+2*2Utavhh{!MFU>Ia1(lv_5mckn8{TF8pB1y{J<y*^&Cy
z@GigLqg5fyWQ{r{U$Ubm8tnIZ|Dlq|-xbu~_|1m)JQX}wXsRLkLc$o4u)9aCN%au)
zoFN>$T0wE&lhaQ3TGg*YeNj6VCkq+6QZ^Plz1uZs{sW}CKguC7nps_G?$kCW=oe7&
z^zDMG6ha04niUt@`7`q+&9F)TU*i@s%fkB?vK~GWo-d<yCG}NgIV}GowQyH6xSqpj
z4lMtY+!?r2ojUJ|CMWOlZ*V6T$nb<=`y^tNi*`US@toMQDef`ZXszwe9#eVZm2vHr
zr0+AUyQZ?r>ubyEpudM2ynN4SN!eMh_D&nsw||nu)B)#`wK!wlGJ2Y-&~Bjb*k^{|
ze#z(e?C4KdD4NpAL<WQ!)y?M9qj)4(eBn%>ug;ZHHWcs_$KP#uzkv8KS9oAaoB2Hk
z`@jxeANWC^7W4}%P{-_9iczllaBAsQGOYmzO<LAR4$ZRpNP;NY9}9Y_WL|5bH4S;E
zFG7i|)|t}gv<6!`mW_pvg=zlYHHn`Gz3$UrC}@tH^ABW@JBLvrGQuJ9A6g*Zx+0rZ
z3VF+i1N3<<&I>Bbd;II7p1NfWNDHd;el{Dm0+z0$YxkytIru=I!uUXPrcF1FcF`v$
zF<z-CSnQQVf7kCt0i`eF|8l~QKyP1UT@m$!-%W}dH_>)FvMW*eB)lN>v#C%J!MxJ1
zS^)Hml;4xd_okG+sh9as9EDET8J>{4<$CJqxao=7<#}+>!yM?;{6Cvb{)7pR#;lRZ
zzHGR{9OC&<=I62<otblI{tw|Yb6-}L28XU5t`WFag?{(^&&gU{$@*v|+LP4!w3{95
z+f?<o2#DTX88mKZF}Mn8X!-cZypK&?>>%N}(mhwW2IxoLsndux+E6Y-;uXmC#nVTg
zXeHmSXDS}=9tWbE*3P`Qc581{vU_Q!p~9O>1Mne{!A8Obw7@EVOwmV~P3CNBu#ZU3
z$ZKFu@n+@Ny2K?%Gj={tVJA0=XG-opXI<|7k9wdlA^VQmMDA#X)|?#G3^tpyYCVY0
zEN{w7awk8_VILC({edTYG#~6RxIj^No}L?>z%F;8nJUX=lL@jFeOqoe2J{v%$eqij
z&jKq;XC2~Q`@FALcQtoJ59!fKByJ^S6u$tyfH|$~0?a3*2Qep7dSc?Vzpp8wP~^LW
z{ZJ-y-;DJ!=yOEXf&#G*9$SdbmOuor4_u)>M7wbZ85g3lenIp>XU;3V_90)Q>Qu9G
zB+M6*vdMV6&%yZbgTTwF1P`tN+!7VoU(lb*poP;Rn%S?#gv+PkEFHy1>3fJ-cb3~m
zTu1~3(5I)iiDPVnFeZ<3y!SXoH+^8?W=wmqbq$C}H{KKd74-V_dJ(ENbC1TGEh;_d
zPY3Mm$GmUfrawG=I>SJsgnR=1zteus^|^M-gHxG&t!c&FBw0A$(58ys7ifC%tZHl{
z=(kG6)zv(IABz;$Hu~D9%+}xO&A`KA)gVygk;FxXGw&5-=MO~TlNDZ3r=aIIZ{|Ne
zHjDevs$LtZ(|6JwaW%4FpV2IPPGau95N+f+ztU=@4r~AoWa+Xe|E(a;4oJN03FtQ*
zd`C(78-6aW;(EueK9zy(t(!@ET&|x91-;fYF-4%KPf#2?9cN6!UtGiN<d^+78~#HW
z$vpl_B=5OOr^fY1pg-9=M7^N$US&Sew(SV1X@5~V6yQ($wav4;=g~DRJ2U8Ua+Y+C
zOMYB2am}XNRYm&xk<=njuU-u64ZzXo+oj}yzT<6$SLMt#J|6{5uRF4gf^1EnkBldo
OFuNs&3d|32tp5jf1uQoJ

literal 0
HcmV?d00001

diff --git a/ee/audit/test_out.txt b/ee/audit/test_out.txt
new file mode 100644
index 0000000000000000000000000000000000000000..c406046d4a5253828013b4c4039604002ee0ab47
GIT binary patch
literal 4768
zcmV;R5?}2VBilxmY@Fntd^C?J<zg#r<K|+cxP!pBGyDg7jIj))fnVBcMAa)IOVjX1
z;hCmxIX>gD&7ha1S-mA9u*_CiC$ZIUNvC>BF?5ecUqY<sg9!lHtQ?SzQ(1t1N`(3K
zxWZc+e72u8KB-Cgq0*^F(I#9~Kx>q~sGm}AsCbB$W!L<=1r7d#OVK*0g2+-%+(ss=
zp)XtS1dw~azCh9NF4~5gc`Rkta7w#vy>9K48%VFkK{j6kW^R3#J7hie$|0|bd6z(j
z=jGARrk#?1`V$Js=<rd?mj}w06#LZv&(n_p&&B`_Di2HvUy-u8;6W?Q$IOoSJagyM
z-Q9!Hdqr_g#xk@EbWaT~%IdRfdGDn{p5H1nRZjuHB0E5i_m`Q1{z572)XGOh0`xS=
zK-Jw+K#1*h`@I4yY3G}Y%=ll9lV3K__M6Szefms5?FF#*R@eCDa-S$*U|oUV->R!L
zuWc)dEvK@rIhsuR+_<TP!kcobaA^vh6!rXop;|>cW1VTsyDFl5s8w%0F0VF8JOMca
zgAlBAdl?>VJTEI%Y(~w(9sSDBtL?-!_sK=P#(FV`HQ0P^O~q=k*DdWJ{Vh7i=4{A+
zu5xC&UH-BhZ*W0;()+7m7_uIphEh^E1*EwJ`-NpNTF>rnynHkXERFCH4<)2ZMtMQx
z^pH83LYcJI{p2Qoe$m|x6(iM^3B)0oZpVvb`Rmo>ZqEo`|3I$Uz{mkc#^Q|rEPE{!
zG2TjqQFJ_W{@Oc!=S_u{0tvZPm=zp>Z4nl@O{Srx!=vc$o(u14;(T>`@ckmMSMK+;
zPbHw*ARWXPAdHXKI_j1svb=_TEb>~K7*VN<-8c~#nljp~ruh@fEl#qMEqBQrP&4Lg
zP<=%#i42V_@T@;CuZM|8esu`pJ$f4U@LEc#5!Zedj}Yf2As?Tv-&ye0i6xViL8Po-
zfUz2v1>ZVS6WX${Fi&{*ilSZa%@mu+$Q&Ym06GK^h%*N_^aS3NbaB>{2PxeGUPq8<
z0ot`RFy%TOPpl=1tU{U>Ej>{eOXz_v6CeB}w5hY=>KyU|g_#@bf*qASF1X}y*a@tl
zsSou=@x~XrQ);KxF&h=Bjz6<Jq_S_UA{^{2?gl1~ODdkz{sGqiOS&RuvEWa#+jH%o
zgS%z|%EME^X~-aJM2o$g5L~r_leW%hr$5X`lKvi`rW9Z1xT|B8Bap|CY#T2%8_6X)
z==|${K|63B<jyS!KIvQ87e1rh=UpFFmoMMrX^wY|pP%l59e$Duq_A*xVZlwR&a5rw
zLT+M%vN!}a3<Y*4V@}WHM5g$@pjc$Jw?rY--gQoEo3~03_gQx5Yz%tTP9uoBW<l37
zD<^6mtUnlub$sr2mnl9i+arwP&Lk#14J?_+GE!{Z@+zsb2&}$30uA4?hWR%(LkIzB
zmhG;Os4psb%t&2P5z82P{}Wd9wtMXujI(4>gXkWJfT|l)M~Y?0S|qMCmv@ZUn~Ued
zseccO)l%ovt&V-|OpWXmG@ocugwshEEJMdfGaiAKKe;CD+TT><pD1G2Uf8wr8@5Ux
z*U%K<&r=I<u~7^|A*xzA6LK^qz8zU)j5MKyI)rY-`BwEJ1R1txr~`M!eI8&5;=dJ~
z(%WwT>CKzyeEj?io+(!=6GtB!#h*MO-Bo&ffljZYwHuQ-t#lXRR^w8*H2+RZVhb6M
zQGmjKTm4#$U2c4O!Pcq1QRf-K)up>>tOnQChFeZslH}O)bHaXqJ^}1{@yl1<Pr)@l
zlZ5VuL<T$Ro=u$e4UJ`GgR=d4K5NVdEFts0PK!1AUP{{!y@Wl9f#hL^w+QReY!w|k
zGf|y|Wpt-Y^TGc%6GfCmLnE;+a#<}%;uNZ{+`$3LgH|0>xF4icXNMiQ8bQCk3IZYz
z2Gq-;eQA*dgj$7Rf$p{Dve0FD9@;W$t$L6+e)H*=wjspx>h*Ts#M}&_a(254Ql=53
z(X`@!<Zh(6{IZhwx**;{P`Opy1@Rl%6q11tTEI`~q{YU(6`J5X<}!TeR2LfYgIEUB
z+Iqn^pua%83=HKO75KG&GhcqmSp3r#p|5u+<21B&n4fVoL!_SRWzI6j?=qCuPj9Pq
z2PR&V+Etr2hv0^<QDupouV%7F)J{IFP);+P8{tGBzj<;lE2qbg>c-?~K<QzCXtjxr
zk=n6iW3221o{g`_eXd5$p2!_!$)sM*b@|HofDFKaz=WV)b&HuAiLdOY!hfO3jVN20
z8z;OJww04I-o_}AD5SvaHsQZNB~yFN7RWrT&y$I7f=6T(LCwVZcglo*7C~=|C|Uj>
z%Z2seNyU78_LE0I4JhCs1HnSxJ5q=?Sn=37L*%bRAr98kqaJjU#0Qs5^-_8a#{bJ&
zV@DJTHPt`MhnE~3Nm93t7JxWKKNzDa+oiy?HeIOBcCS~|Sl7R=PY25V0*tR?C3sRw
z_KAP0HcE2<d4|bk#*<<DJbzByG$U_vLj0#(W0D<L_GfvR9|JaV1P8Ygr5mf7soiQH
z9}!p{oEI*9UMj_|%8RrYuViAA@koczw1Iy(50<onbH~GkMNK!(iC6f7AS<=$9_&_W
zq9>>0c{mzMzRWH!);w0M^;x`igxex|$?;}6Ty>;1KU7R{21=t`c2;l?t?$Elj*@%8
zL`6FC_0n$JZX^<*rwVN$;i~Lq`aVDR6=%cf$Ec9aiQWtbwMckuTRj4<CJ!!br+Bh}
zCds<M$r(Y)!M!qh--)q(h-9%S-j#LF-=tb=><6&B&AAgspp;8uU!{RhItFHlJkX%@
zzGK1Qo~T0y3s$BAZb~>=uyHL87ho1dLI3XlsKDk<jvyNmh0mN$Yq-dRz;w))szkq#
zOrIi`WhGdLqKjK64Mo39z<zrlLWOEq);y}kqQo&)$Ry^=HVx8yqE|@g9|D2$=u>DM
z-_Vc}wh7H4Ex#&mK%p9Nv6NuN>34=6Sr7I48<7LQ7oZ91dhq9Tb4HcgcDYX|{{$AO
zdO^v=yG5h7c;i4Q<0wdAC3OBb2zS1ISw>c%8`zfk&J<xnuaJd5ObWy>7Oz?x7|=g_
zEr_2|<B)v@9rcq6&#k@1z~i=mn<1@<aP>|fqB!$GGL!lMnq_vHW*<&AYU5Fxc3cgv
z5MWyIPoIho`h1hB_G;oHMn0I<O$C09r+G9FB-3@9O7w&s0M}YSCGOQm3Dqq|bd_nt
zHlT%G-4WCo%=YOrYP_(p+9QU8Bbor^8sqhJ)SDB{NiH!U;EQyb+{-a75G3v8h3(Aq
z<d!dvjID%J!cDg-ue9<eSPTh8uvWMhQd=_MRJX`EVm|EqIXoL-jun|Po)Q2@>Epge
zdd~&xwp@ME2#2hrZaRZ2hg;9w-GxFQ$kTDR)|LlY0k-+wA&^CCH2AbS6?yZa1zX)B
z8Twz^=j)F$?xtuF&}ViMJo5B)Nbz8hY@2#!qer?B;Kp7>^CAuQlQTtlQJ=^ee~HoJ
z_VbnxBYd4(>L3c1h>7Y|HfvSKfPB<`{U_7-`9v#e-X&%B0}o5Nr$y#46Df^cxI|S?
zpE8y-QYdn*7sRiZEpU-;&9c0Aii1bcIRBs>d&)^4kWq+rfHhZe_O09zERd(B<QLJI
z`gc9^{BCw9D?eTcy{yG51%At+mH|K)#JN~Rx6{^hAvr7_xwi*KjvXMoJjzJ}imQ!0
z4^OdKo(iEf!m8MIlpL?#sQrOAEGQr+FSg@Lv#M4K$Pl9Ib4loDDGlhFnj6y;Wb0#T
zsTu7HrWKMSC~pM+ud)N}M;b+QC8&C5%61mY*()=oIH;?ojY9_@?C^rS3F!b+)AT7M
z>SmLY=Ig-56~z`e53tVs^SCThPRY&02;I)92oYJrj^y%4T3(Ep4Z63vDsyCr>XGhA
z{oj<c(qsDghXK{|1{Nxj`*~k2xy#cvJ>eP(;|_pZ^(;ZGCbB4W&wcSJ-6nqE&jaMi
zoC~C))dgN}ZJ>ru7|ZLbul_)UKq!D}-7GM!z5Bn!zx*#LA=d9q)2N%W_G)oY6QvLw
z{pS#FLhsO0ED0hrM?rit@*AqN$N&T0uwN`@NF^lO+#PvfPqn5F<5;aTuYOFEz^I2r
z?}hZnG;M0EpW~W&k+BvH7iTbT{F-J#VBE!L|HuI=a`=$$;^VxuEU_;)5yQjzGJ5B`
z<|5m1PLa=aytt}tsXe)W%Mo9C$2pMA>o+3m-U_xvx%jL1ev~W*^Vo?|3D6(xyG-{*
zvYbhD9uK}F$K?vrXV0oEA?!4$m;(kTfns}{*Q`hLJn7a~(n3GV1y}S&f8vyIuW?=%
z4m7C-acAL%uV5d~JT4W|#J|mC!kcM+gUTCYN6S*1xD6DH1IlhKhlieU55yDPgqL%>
z;jG%1O)F_MZXW~>euFC}hnqLb{F6zEq?XT^&577;g}@tp-@Os@hu;sKrcyT<B;>n<
zAiaG<4L>>T)p$ikoCeIiw#+M^x8J2sKO{MqK}sTTSy_A5FEDctEbm^lfW{s`PC78b
z$`69xtU%sPNvN0w=A{Zu)*PsSL=OCK_+9E4CBBV*^a{O&RRZ#x2JB;<>;_-|q6;%}
z@Wa6)@J#rcvk$t(D_TcqB{_@(oo%0ub>4?Op-4*w+Di{s#dO|PbZn*cF8-<Dl691O
z`DV!+gw#H{o9cMXr*C=STAK~4szYP`Dbn2^v7Z~O4#HE8K8gV!61VO2V-s#yL;Y_+
zMCuLe70leC=S|7u8IFMFe0aG=3GZ%vmhTieIh<>Fgvkw5z&J=0_DPDF`YDSQGN}9F
zyUc*YQ(G#?enuS;q-e2{Sp*?m1<`g@UN_4S|C-RshagYW(!@dN#7ElDt{6<SqB!uT
zpqwAh4|*Awmn*a{eRLuW;^|tlL#xhv?ydjr{I&2s<UzaZ<A(i!762ds7Wp+qS^Z`z
z=sgzeR}YFghYVeK%73x&^{pGn33+EJ?o|rDMB<L!h=YitZQb5_8E6d8cHg%*^CaXD
zXX4s1YkLs*0HaE!Q~JtJXI^r=HR2v7qAz_8UI>X4l5HS7!4_z9r^BPL;v1sSuqNqV
z!nsZ-v+Z^1iJ$u|c0lvklZY>ufEtN}Ygrl?x7Sh%2~tpiM0SHvz~$_@1XZ2#)g3B(
zS1Ud|zmMx+|03UL$^=64G3r=+^#OMz<g9MT5~`-B>V`-6>{)O3L1WV-J@74Y3PW85
zKgIh#xcZF!OhPo2MnZeRnyhbsmknyy_1wmr$|pWn2$_fW85y&x$sIMOf+s2SceBOt
zZL@YKdZ9<KBf~!1kD;FJtdLsOhn4OH`tftkWxQ5qJrxPuw-F?=r!*a9V65Jij04_P
zt$mdIj$8D9g6ob`DKYBe0#okB>rmFnO4%vGmSFOwVK2ph1DY%dUbOGD<58_{<6Ixe
zD}1Za%b%hbF6;U*M$~vb!{~ZQ>f!z!QrpfGi#$oOs5Vl6%|az$WN8zz^HpuC7b+2M
z@vhFU9tqFfv}x05edEq$;l!^f0JMs-hPLq&YE;ko>i&UhDegC(gH1?|ikw=DW!`Zu
zo42a=dG-hWDCUb<r98neqAZ&A^xWQCp#KT|nQheSH7r00!%_meZ!wFwv+7ED?J+i;
z*Sv=eqFDm+kMamQ#R#Ls<N;7a@20ko%Xy<w&QmjCDA4UWWzrRRM8J2WHX|9bl^G2S
zk+^Y#WosPcdrU#a_uzYnj=G<4|3oevP6;G$5p^dZxDr7zytwm9mrdcK)k7vWRF`Ex
zI+D_v?n+VDhFOR)<2fjUO8uewfCtpjV`-&dSaVkG^kT#L<2qANhPX@1i=x&SXxDRz
zN?Ra$E#<~F{sJK5NF@<u@D6W)zb-*1!!Q5m`0nrQOTzGRf>+72K@8o(XdVc|@9>oG
z$p{B{`!{!pGecb)j^JFW>m_XSpNW%WQ8FwNXn+5VqRJi#)mr`Yn91!zn0c>(JYfyA
zEyzXmA8JZ+GJ`-p{Qo2<xizRb#U7mZzdIdH$(C2(bL@lf89f<Pnfzx+M{kukc>(5O
zspkP&Za0KO*Qe|f!m44on>E!rV70Zq4=uWok%jRPrRW|u)LK4?_H_2)I_s|kWnlcA
zw30cki8=M8b(SaY`DN({dCt=rIdF4GYDwnj#U$pJN$$_tC(}uA(CUj{H4M>B^VCgt
zQi;@Ipe|?8a7c=K7W>4l99=ANV?`pdqKc1qg4vg@DC7duP}G#*>me**f=EgkY8r|0
ukj%(<f9$;=^I@cmJd|=<c$63i9^}*YZc;h8gyQ;LcloEV3yBA*#QS@J9Y*K?

literal 0
HcmV?d00001

diff --git a/go.mod b/go.mod
index b77791e7fe5..63b6afaa57e 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.12
 // replace github.com/dgraph-io/badger/v2 => /home/mrjn/go/src/github.com/dgraph-io/badger
 // replace github.com/dgraph-io/ristretto => /home/mrjn/go/src/github.com/dgraph-io/ristretto
 
+replace gopkg.in/natefinch/lumberjack.v2 => /home/amanbansal/go/src/gopkg.in/natefinch/lumberjack
 require (
 	contrib.go.opencensus.io/exporter/jaeger v0.1.0
 	contrib.go.opencensus.io/exporter/prometheus v0.1.0
diff --git a/x/log_writer.go b/x/log_writer.go
new file mode 100644
index 00000000000..a3dba65170e
--- /dev/null
+++ b/x/log_writer.go
@@ -0,0 +1,120 @@
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package x
+
+import (
+	"compress/gzip"
+	"encoding/binary"
+	"fmt"
+	"github.com/dgraph-io/badger/v2/y"
+	"io"
+	"os"
+	"sync"
+)
+
+var _ io.WriteCloser = (*LogWriter)(nil)
+
+type LogWriter struct {
+	FileName      string
+	MaxSize       int64
+	MaxAge        int // number of days
+	Compress      bool
+	EncryptionKey []byte
+	size          int64
+	file          io.Writer
+	mu            sync.Mutex
+
+	millCh    chan bool
+	startMill sync.Once
+}
+
+func (l *LogWriter) Write(p []byte) (int, error) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	if l.file == nil {
+		// open file
+	}
+
+	if l.size+int64(len(p)) >= l.MaxSize {
+		// rotate
+	}
+
+	n, err := l.file.Write(p)
+	l.size = l.size + int64(n)
+	return n, err
+}
+
+
+func (l *LogWriter) Close() error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if l.file == nil {
+		return nil
+	}
+	l.file = nil
+	return nil
+}
+
+func encrypt(key, src []byte) ([]byte, error) {
+	iv, err := y.GenerateIV()
+	if err != nil {
+		return nil, err
+	}
+	allocate, err := y.XORBlockAllocate(src, key, iv)
+	if err != nil {
+		return nil, err
+	}
+
+	allocate = append(allocate, iv...)
+	var lenCrcBuf [4]byte
+	binary.BigEndian.PutUint32(lenCrcBuf[0:4], uint32(len(allocate)))
+	allocate = append(lenCrcBuf[:], allocate...)
+	return allocate, nil
+}
+
+func compress(src string) error {
+	f, err := os.Open(src)
+	if err != nil {
+		return fmt.Errorf("failed to open log file: %v", err)
+	}
+
+	defer f.Close()
+	gzf, err := os.OpenFile(src+".gz", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
+	if err != nil {
+		return fmt.Errorf("failed to open compressed log file: %v", err)
+	}
+
+	defer gzf.Close()
+	gz := gzip.NewWriter(gzf)
+	defer gz.Close()
+	if _, err := io.Copy(gz, f); err != nil {
+		os.Remove(src + ".gz")
+		err = fmt.Errorf("failed to compress log file: %v", err)
+		return err
+	}
+
+	if err := os.Remove(src); err != nil {
+		return err
+	}
+	return nil
+}
+
+func manageOldLogs() {
+	// delete old log files.
+	// compress log files
+}
\ No newline at end of file
diff --git a/x/logger.go b/x/logger.go
index 55f60be5b74..014f92805dc 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -25,7 +25,7 @@ import (
 	"gopkg.in/natefinch/lumberjack.v2"
 )
 
-func InitLogger(dir string, filename string) (*Logger, error) {
+func InitLogger(dir string, filename string, key []byte) (*Logger, error) {
 	if err := os.MkdirAll(dir, 0700); err != nil {
 		return nil, err
 	}
@@ -34,11 +34,15 @@ func InitLogger(dir string, filename string) (*Logger, error) {
 		return nil, err
 	}
 	getWriterSyncer := func() zapcore.WriteSyncer {
-		return zapcore.AddSync(&lumberjack.Logger{
+		w := &lumberjack.Logger{
 			Filename: path,
 			MaxSize:  100,
 			MaxAge:   30,
-		})
+		}
+		if key != nil {
+			w.EncryptionKey = key
+		}
+		return zapcore.AddSync(w)
 	}
 	core := zapcore.NewCore(zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
 		getWriterSyncer(), zap.DebugLevel)

From 40e019429143924252deb6f1fcfc2a9c1cc96862 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Tue, 19 Jan 2021 20:35:09 +0530
Subject: [PATCH 17/31] adding logwriter to handle encryption and everything

---
 dgraph/cmd/alpha/run.go                 |  17 +-
 dgraph/cmd/zero/raft.go                 |   8 +-
 dgraph/cmd/zero/run.go                  |  32 +--
 ee/audit/audit_ee.go                    |  33 ++-
 ee/audit/encryption_test.go             | 153 --------------
 ee/audit/run_ee.go                      |  28 +--
 ee/audit/test_encryption_chunks_cbc.txt | Bin 4800 -> 0 bytes
 ee/audit/test_out.txt                   | Bin 4768 -> 0 bytes
 ee/utils_ee.go                          |   2 +-
 go.mod                                  |   2 -
 go.sum                                  |   2 -
 worker/config.go                        |  21 +-
 x/config.go                             |   3 +
 x/flags.go                              |   4 +
 x/log_writer.go                         | 262 +++++++++++++++++++++---
 x/log_writer_test.go                    | 157 ++++++++++++++
 x/logger.go                             |  13 +-
 17 files changed, 496 insertions(+), 241 deletions(-)
 delete mode 100644 ee/audit/encryption_test.go
 delete mode 100644 ee/audit/test_encryption_chunks_cbc.txt
 delete mode 100644 ee/audit/test_out.txt
 create mode 100644 x/log_writer_test.go

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 0ab23a970a7..4b9e6208982 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -198,8 +198,14 @@ they form a Raft group and provide synchronous replication.
 		`Cache percentages summing up to 100 for various caches (FORMAT:
 		PostingListCache,PstoreBlockCache,PstoreIndexCache,WAL).`)
 
-	flag.Bool("audit_enabled", false, "Set to true to enable audit logs.")
-	flag.String("audit_dir", "aa", "Set path to directory where to save the audit logs.")
+	flag.String("audit", "",
+		`Various audit options. 
+	enable=true/false enables the audit logs (default behaviour is false).
+	dir=/path/to/audits to define the path where to store the audit logs.
+	compress=true/false to enabled the compression of old audit logs (default behaviour is false).
+	encrypt_file=enc/key/file enables the audit log encryption with the key path provided with the
+	flag. 
+	Sample flag could look like --audit dir=aa;encrypt_file=/filepath;compress=true`)
 
 	// TLS configurations
 	x.RegisterServerTLSFlags(flag)
@@ -612,10 +618,10 @@ func run() {
 
 		MutationsMode: worker.AllowMutations,
 		AuthToken:     Alpha.Conf.GetString("auth_token"),
-		AuditEnabled:  Alpha.Conf.GetBool("audit_enabled"),
-		AuditDir:      Alpha.Conf.GetString("audit_dir"),
+		Audit:      Alpha.Conf.GetString("audit"),
 	}
 
+	x.CheckFlag(opts.Audit, "dir", "compress", "encrypt-file")
 	secretFile := Alpha.Conf.GetString("acl_secret_file")
 	if secretFile != "" {
 		hmacSecret, err := ioutil.ReadFile(secretFile)
@@ -676,6 +682,7 @@ func run() {
 		TLSClientConfig:      tlsClientConf,
 		TLSServerConfig:      tlsServerConf,
 		HmacSecret:           opts.HmacSecret,
+		Audit:                opts.Audit,
 	}
 	x.WorkerConfig.Parse(Alpha.Conf)
 	x.CheckFlag(x.WorkerConfig.Raft, "group", "idx", "learner")
@@ -799,7 +806,7 @@ func run() {
 	adminCloser := z.NewCloser(1)
 
 	// Audit is enterprise feature.
-	audit.InitAuditorIfNecessary(Alpha.Conf, worker.EnterpriseEnabled)
+	audit.InitAuditorIfNecessary(opts.Audit, worker.EnterpriseEnabled)
 
 	setupServer(adminCloser)
 	glog.Infoln("GRPC and HTTP stopped.")
diff --git a/dgraph/cmd/zero/raft.go b/dgraph/cmd/zero/raft.go
index 94749e92384..a4155e55e86 100644
--- a/dgraph/cmd/zero/raft.go
+++ b/dgraph/cmd/zero/raft.go
@@ -402,8 +402,12 @@ func (n *node) applyProposal(e raftpb.Entry) (uint64, error) {
 		// Check expiry and set enabled accordingly.
 		expiry := time.Unix(state.License.ExpiryTs, 0).UTC()
 		state.License.Enabled = time.Now().UTC().Before(expiry)
-		if state.License.Enabled && opts.auditEnabled {
-			audit.InitAuditor(opts.auditDir, opts.encryptionKey)
+		if state.License.Enabled && len(opts.audit) > 0 {
+			if encKey, err := audit.ReadAuditEncKey(opts.audit); err != nil {
+				glog.Errorf("error while reading encryption file %+v", err)
+			} else {
+				audit.InitAuditor(x.GetFlagString(opts.audit, "dir"), encKey)
+			}
 		}
 	}
 	if p.Snapshot != nil {
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index d2910e57078..f15846451cb 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -57,9 +57,7 @@ type options struct {
 	w                 string
 	rebalanceInterval time.Duration
 	tlsClientConfig   *tls.Config
-	auditEnabled      bool
-	auditDir          string
-	encryptionKey     []byte
+	audit             string
 }
 
 var opts options
@@ -102,10 +100,14 @@ instances to achieve high-availability.
 	flag.Duration("rebalance_interval", 8*time.Minute, "Interval for trying a predicate move.")
 	flag.String("enterprise_license", "", "Path to the enterprise license file.")
 
-	flag.Bool("audit_enabled", false, "Set to true to enable audit logs.")
-	flag.String("audit_dir", "za", "Set path to directory where to save the audit logs.")
+	flag.String("audit", "",
+		`Various audit options.
+	dir=/path/to/audits to define the path where to store the audit logs.
+	compress=true/false to enabled the compression of old audit logs (default behaviour is false).
+	encrypt_file=enc/key/file enables the audit log encryption with the key path provided with the
+	flag. 
+	Sample flag could look like --audit dir=aa;encrypt_file=/filepath;compress=true`)
 
-	enc.RegisterFlags(flag)
 	// TLS configurations
 	x.RegisterServerTLSFlags(flag)
 }
@@ -197,7 +199,6 @@ func run() {
 	tlsConf, err := x.LoadClientTLSConfigForInternalPort(Zero.Conf)
 	x.Check(err)
 
-	key, err := enc.ReadKey(Zero.Conf)
 	opts = options{
 		bindall:           Zero.Conf.GetBool("bindall"),
 		portOffset:        Zero.Conf.GetInt("port_offset"),
@@ -207,13 +208,12 @@ func run() {
 		w:                 Zero.Conf.GetString("wal"),
 		rebalanceInterval: Zero.Conf.GetDuration("rebalance_interval"),
 		tlsClientConfig:   tlsConf,
-		auditEnabled:      Zero.Conf.GetBool("audit_enabled"),
-		auditDir:          Zero.Conf.GetString("audit_dir"),
-		encryptionKey:     key,
+		audit:             Zero.Conf.GetString("audit"),
 	}
 	glog.Infof("Setting Config to: %+v", opts)
 	x.WorkerConfig.Parse(Zero.Conf)
 	x.CheckFlag(opts.raftOpts, "idx", "learner")
+	x.CheckFlag(opts.audit, "dir", "compress", "encrypt-file")
 
 	if !enc.EeBuild && Zero.Conf.GetString("enterprise_license") != "" {
 		log.Fatalf("ERROR: enterprise_license option cannot be applied to OSS builds. ")
@@ -233,9 +233,15 @@ func run() {
 
 	wd, err := filepath.Abs(opts.w)
 	x.Check(err)
-	ad, err := filepath.Abs(opts.auditDir)
-	x.Check(err)
-	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", opts.auditDir)
+	if len(opts.audit) > 0 {
+		dir := x.GetFlagString(opts.audit, "dir")
+		if len(dir) == 0 {
+			glog.Fatal("audit flag is provided but dir is not specified")
+		}
+		ad, err := filepath.Abs(dir)
+		x.Check(err)
+		x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", dir)
+	}
 
 	if opts.rebalanceInterval <= 0 {
 		log.Fatalf("ERROR: Rebalance interval must be greater than zero. Found: %d",
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index 5d59a45436d..9ef5e0962e1 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -15,7 +15,6 @@ package audit
 import (
 	"context"
 	"fmt"
-	"github.com/dgraph-io/dgraph/ee/enc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -23,12 +22,12 @@ import (
 	"google.golang.org/grpc/status"
 	"io/ioutil"
 	"net/http"
+	"path/filepath"
 	"sync/atomic"
 	"time"
 
 	"github.com/dgraph-io/dgraph/x"
 	"github.com/golang/glog"
-	"github.com/spf13/viper"
 )
 
 var auditEnabled uint32
@@ -59,23 +58,39 @@ type auditLogger struct {
 	tick *time.Ticker
 }
 
+func ReadAuditEncKey(conf string) ([]byte, error) {
+	encFile := x.GetFlagString(conf, "encrypt-file")
+	if encFile == "" {
+		return nil, nil
+	}
+	path, err := filepath.Abs(encFile)
+	if err != nil {
+		return nil, err
+	}
+	encKey, err := ioutil.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+	return encKey, nil
+}
+
 // InitAuditorIfNecessary accepts conf and enterprise edition check function.
 // This method keep tracks whether cluster is part of enterprise edition or not.
 // It pools eeEnabled function every five minutes to check if the license is still valid or not.
-func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
-	if !conf.GetBool("audit_enabled") {
+func InitAuditorIfNecessary(conf string, eeEnabled func() bool) {
+	if conf == "" {
 		return
 	}
-	key, err := enc.ReadKey(conf)
+	encKey, err := ReadAuditEncKey(conf)
 	if err != nil {
+		glog.Errorf("error while reading encryption file", err)
 		return
 	}
 	if eeEnabled() {
-
-		InitAuditor(conf.GetString("audit_dir"), key)
+		InitAuditor(x.GetFlagString(conf, "dir"), encKey)
 	}
 	auditor.tick = time.NewTicker(time.Minute * 5)
-	go trackIfEEValid(eeEnabled, conf.GetString("audit_dir"), key)
+	go trackIfEEValid(x.GetFlagString(conf, "dir"), encKey, eeEnabled)
 }
 
 // InitAuditor initializes the auditor.
@@ -99,7 +114,7 @@ func initlog(dir string, key []byte) *x.Logger {
 // trackIfEEValid tracks enterprise license of the cluster.
 // Right now alpha doesn't know about the enterprise/licence.
 // That's why we needed to track if the current node is part of enterprise edition cluster
-func trackIfEEValid(eeEnabledFunc func() bool, dir string, key []byte) {
+func trackIfEEValid(dir string, key []byte, eeEnabledFunc func() bool) {
 	for {
 		select {
 		case <-auditor.tick.C:
diff --git a/ee/audit/encryption_test.go b/ee/audit/encryption_test.go
deleted file mode 100644
index e5b529b8596..00000000000
--- a/ee/audit/encryption_test.go
+++ /dev/null
@@ -1,153 +0,0 @@
-/*
- * Copyright 2017-2021 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.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package audit
-
-import (
-	"bytes"
-	"crypto/aes"
-	"crypto/cipher"
-	"crypto/rand"
-	"fmt"
-	"github.com/stretchr/testify/require"
-	"io"
-	"io/ioutil"
-	"os"
-	"testing"
-)
-
-var key = []byte("12345678901234561234567890123456")
-var entry = []byte("abcdefghijklmnopqrstuvwxyz")
-
-func TestEncryptionCBC(t *testing.T) {
-	file, err := os.OpenFile("test_encryption_chunks_cbc.txt", os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
-		0600)
-	require.Nil(t, err)
-	salt := make([]byte, 8) // Generate an 8 byte salt
-	rand.Read(salt)
-	data := make([]byte, len(entry) + aes.BlockSize)
-	copy(data[0:], "Salted__")
-	copy(data[8:], salt)
-
-	for i := 0; i<100; i ++ {
-		copy(data[aes.BlockSize:], entry)
-		padded, err := pkcs7Pad(data, aes.BlockSize)
-		c, err := aes.NewCipher(key)
-		require.Nil(t, err)
-		iv := make([]byte, aes.BlockSize)
-		rand.Read(iv)
-		cbc := cipher.NewCBCEncrypter(c, iv)
-		cbc.CryptBlocks(padded[aes.BlockSize:], padded[aes.BlockSize:])
-		file.Write(padded)
-	}
-
-	file.Close()
-}
-
-// pkcs7Pad appends padding.
-func pkcs7Pad(data []byte, blocklen int) ([]byte, error) {
-	if blocklen <= 0 {
-		return nil, fmt.Errorf("invalid blocklen %d", blocklen)
-	}
-	padlen := 1
-	for ((len(data) + padlen) % blocklen) != 0 {
-		padlen++
-	}
-
-	pad := bytes.Repeat([]byte{byte(padlen)}, padlen)
-	return append(data, pad...), nil
-}
-
-func TestEncryptionInChunks(t *testing.T) {
-	file, err := os.OpenFile("test_encryption_chunks.txt", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
-	require.Nil(t, err)
-	iv := make([]byte, aes.BlockSize)
-	rand.Read(iv)
-
-	block, err := aes.NewCipher(key)
-	require.Nil(t, err)
-
-	stream := cipher.NewCTR(block, iv)
-	for i:=0; i< 100; i++ {
-		e := []byte("abcdefghijklmnopqrstuvwxyz")
-		stream.XORKeyStream(e, e)
-		file.Write(e)
-	}
-	file.Write(iv)
-}
-
-func TestDecryptionInChunks(t *testing.T) {
-	cipherText, err := os.Open("test_encryption_chunks.txt")
-	require.Nil(t, err)
-	iv := make([]byte, aes.BlockSize)
-	stat, _ := cipherText.Stat()
-	msg := stat.Size() - int64(len(iv))
-	_, err = cipherText.ReadAt(iv, msg)
-	require.Nil(t, err)
-	outfile, err := os.OpenFile("test_encryption_chunks_out.txt",
-		os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
-		0600)
-	require.Nil(t, err)
-
-	block, err := aes.NewCipher(key)
-	require.Nil(t, err)
-
-	buf := make([]byte, 26)
-	stream := cipher.NewCTR(block, iv)
-
-	for {
-		n, err := cipherText.Read(buf)
-		if n > 0 {
-			if msg == 0 {
-				break
-			}
-			msg = msg - int64(n)
-			stream.XORKeyStream(buf, buf)
-			outfile.Write(buf)
-		}
-
-		if err == io.EOF { break }
-		require.Nil(t, err)
-	}
-}
-
-func TestEncryptionAtOnce(t *testing.T) {
-	file, err := os.OpenFile("test_encryption_once.txt", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
-	require.Nil(t, err)
-	block, err := aes.NewCipher(key)
-	require.Nil(t, err)
-	iv := make([]byte, aes.BlockSize)
-	rand.Read(iv)
-	gcmCipher := cipher.NewCBCEncrypter(block, iv)
-	repeat := bytes.Repeat(entry, 100)
-
-	gcmCipher.CryptBlocks(repeat, repeat)
-	file.Write(repeat)
-}
-
-func TestDecryptionAtOnce(t *testing.T) {
-	cipherText, err := ioutil.ReadFile("test_encryption_once.txt")
-	require.Nil(t, err)
-	block, err := aes.NewCipher(key)
-	require.Nil(t, err)
-	gcm, err := cipher.NewGCM(block)
-	require.Nil(t, err)
-	nonce := cipherText[:gcm.NonceSize()]
-	cipherText = cipherText[gcm.NonceSize():]
-	plaintext, err := gcm.Open(nil, nonce, cipherText, nil)
-	require.Nil(t, err)
-	fmt.Println(string(plaintext))
-}
\ No newline at end of file
diff --git a/ee/audit/run_ee.go b/ee/audit/run_ee.go
index 1a4fab90fc6..d3411f47b9c 100644
--- a/ee/audit/run_ee.go
+++ b/ee/audit/run_ee.go
@@ -88,29 +88,29 @@ func run() error {
 	block, err := aes.NewCipher(key)
 	stat, err := os.Stat(decryptCmd.Conf.GetString("in"))
 	x.Check(err)
+	iv := make([]byte, aes.BlockSize)
+	_, err = file.ReadAt(iv, 0)
+	x.Check(err)
 
-	var iterator int64 = 0
+	var iterator int64 = 16
 	for {
-		if iterator >= stat.Size() {
-			break
-		}
-		l := make([]byte, 4)
-		_, err := file.ReadAt(l, iterator)
-		x.Check(err)
-		iterator = iterator + 4
-		content := make([]byte, binary.BigEndian.Uint32(l))
+		content := make([]byte, binary.BigEndian.Uint32(iv[12:]))
 		_, err = file.ReadAt(content, iterator)
 		x.Check(err)
 
-		iterator = iterator + int64(binary.BigEndian.Uint32(l))
-		iv := make([]byte, aes.BlockSize)
-		_, err = file.ReadAt(iv, iterator)
-		iterator = iterator + aes.BlockSize
-
+		iterator = iterator + int64(binary.BigEndian.Uint32(iv[12:]))
 		stream := cipher.NewCTR(block, iv)
 		stream.XORKeyStream(content, content)
 		_, err = outfile.Write(content)
 		x.Check(err)
+
+		// if its the end of data. finish encoding
+		if iterator >= stat.Size() {
+			break
+		}
+		_, err = file.ReadAt(iv[12:], iterator)
+		x.Check(err)
+		iterator = iterator + 4
 	}
 	return nil
 }
diff --git a/ee/audit/test_encryption_chunks_cbc.txt b/ee/audit/test_encryption_chunks_cbc.txt
deleted file mode 100644
index 1330a96edd168ec59b68f576cd100845c2a5f13f..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 4800
zcmZYBcRUpSAHeZCLe`l@GNR)$Dw%OM8Re{$QQ@q}o=IO6-5tr^d;VPZI(wc&=4G6c
z5RPQaip1~#kN=;K*W>efd>)_A`(t9`;$v@TWd(14a+20nAedI(>U^VUyF=8rDa}pu
z;tooZ6zw9QQ$H`BeS>AyM8yQ~UUPI?zlBa;Yrejor22HG?3z@n>FlRBj7VutbGRzf
z8uauqFQnw+hF@Wy;!jx=3!NcZ`(uku3`#>hbKwLF836sD_rRrFqxNB_jaY%udBY*&
zlCMWyu#zsqdBK8C$?^iwzaM5E-M#4uaWoP>*1Xi5Bp<Btji$shV_07!(aO|P4D@2p
zT%>aC|JYMnr5mL_s7sah1#f8{sn_k?q57Hx-<|}0L`rz&G2`-nsr;>J_K%9!AH04}
zk3Aq1v&4Ski15En27Lfu&2sxr$q0wmVfw?-xfFy8{P^wvXf~WK!^ZmgRO~^o?(^TL
z^A0;hGJGq42Wo60iX+>DOz5VJAg`*}Q*%Y0f_^5=+2;>q+UY7!u)yL<jPvVtx1U!S
z)~8$W^<J`UFEOA;_6FMa){bOrWZB~C?CK5$HP5+Z3sb`{rQuHDK0KPBe?2w5>?wV|
zM6Q|giF4ddOg>}vC7Y?ngr=<gB5ao>5A-;#AM=T|93EK%oox}w%9pLGuKpjTykGdh
zodP07znodW#GtvVazL{d^}-P|ejhc<m^t$^L3hRUyIime!lM2=*xw*kzpo{dX&swI
zi(Zymo?D2G?#DF~T{t(K^a%IeWI#U=v=Me|H#o%O=O~rnfI{wse7><{Ose`rb|ol*
zY2nO%H>nvO?z%o)p>w3iFh$=SoTt9KM2ULi)-_Kiv|zt5g8ksq3R{!}-&9~Q&Xz`+
z!08+p*(e&3-h>o)(0;c&U;uiv*&iJb_D|;DAm&Gsn2WbGg@n^>y7CM$d)2-Br87OC
z-;qEFH{D}*WuJ7@JKj|l=PxSX{}$R2nz&_4gSXY(0DU~q2g-*&+K&As7FtZ%&da&6
znY4KUMZ2J0j@E_Ns0Pr(U#+hwgx&x3h?`h%96wJo=w55~*s0B1Z*e!RC2q8Xo?Umr
zk>XI!^VVgz0$-@~U6-ppI0#4XjP=8G-rAcafj&P662BoLt6JvT@CG%K%tOt+X%SgP
z3S{pm6g1z{_60p_+q}l69~Y8+{sqdprtF7wi7-ERjXw4oj(=3RIXw~dW+5&_+PcF4
z-{)-d@!z$IsD`3splHGHC-~(0j|;RA&};qCv0g*E6DC&E{Z(m*Xptpv+BQSf==%%W
z9g~)NmY`23=a}$wBC!dBm6z>Tm*Fy9#$o{$86HhZm?*3kLn-KISex*PN0b$VxAG3{
zZ3t$>#<!@8liazwOSp#5tpqR7quI$}B%*VYRh~-Kq(HoulZ(;A!Jh+f5@t>QgShI4
zfnFx#ci=mreUtsB^2W`_s;;HHE8WUh+<eVOTzLzBsBD5hhi0sWo>lsi^HLqWEmd9N
z3f=}e>jC8@DMco=mqV>U{}Pr*LzN)E8fz9G*GKH8G=!4GNC+E)iJpb>pkF5vpvQb3
zG_thV(75pTl3kbF0|%$Mz3*uQe{Z5|AV2?FgcXB+Aj{qc&rtKt%svIf!g$a}xXF7o
z%q{he=`f6Vr^~Gu^ps?*qB?VO(2FkLAgR{&<58LmJ<?5<wMGWZ+^q;PVbH7cIRDkp
zdC_UW+k_DqXIfnf<jh%;F#FLGJDYWhc0K^~n)LQ^^n;Hy-drtwZFl)RTfd)s))U%q
z*goUbrOXdc-+`W?Dwi*Fi?D6|nzGG5S+}fMef7v*ZE-jvw)=m_h>#V~k0!;6r#x1*
z7tS8H3cM9&UtgO%&(VpNvlx!8(dC@#1U*^5SP6?!owf7WOX~H0hF0I3MtxfEw^P@4
zbZO@hK6AesSS0%e+`SYRU*6~(GaN(gL=(CFQD0OtSg(9EFM?Ts{kvYIcoSIVf~{Zu
z@Hehg<V}D1?tyiUn4VGUTm_?pbkN^JSZpJwS5wAskLLwMn{{1umxj7x)5HUd_xf3i
zpN)b(@k(QNdiK!7jW<UWBMA!Ur=l$lbr{M^nG~tQ6}Ivg^!l{ay=)gVA!^W*KFa2&
znh1}Lo7}=&`^i-v-tjnQ9OxZ&<9s!Te&Z39ntdk%fBA4*k_LX{i^mtD6s>ejx?Mp3
z-Zjbn&@nFc5vc|9ee7$g=O=H@n;86fj=H>v4Wv~B^yLjH!$&D65BcRgS)Wk$q!LWd
z!OsOehxX&<?Be@x+yXr-g@&q)EE(HWmpvVPSi|u=XA62D2Dvnxd0X%JI42PFdD@F*
z!9{xGZ$mF%IFYv64k2DZbQ>#iXEi-#GU!2gfPPLyypWgTw<ZPO4ScRPwzQg4+psX_
zVJH>)ZT~CUkr(uoYT}8LboOh~R^&LL5qh^s`Q4!0QFI0-s=|@Q<3a}LF*(O-qU<u2
zr!Gu-)qKlDp4f9NE{68-BQF`z18D4-`S>Ka?Z(t5IegPTq?N`0hUkRhBL<D+_E_27
z;EbD=rbb|2M{m3Y^Cz#q6c@xrRG4U0)*TOWZ)!M{wbO{(EQ<~Xy;x>ab=)0vd|U#|
zyFRdc{SwxFkB9#;dG>(UNGRfT81zw|SUj6+(oB??F0}>HaL!RUmXkH-Ix*~XHWVhd
zg9m-EF)AOkpp3%@h;E4+2*2Utavhh{!MFU>Ia1(lv_5mckn8{TF8pB1y{J<y*^&Cy
z@GigLqg5fyWQ{r{U$Ubm8tnIZ|Dlq|-xbu~_|1m)JQX}wXsRLkLc$o4u)9aCN%au)
zoFN>$T0wE&lhaQ3TGg*YeNj6VCkq+6QZ^Plz1uZs{sW}CKguC7nps_G?$kCW=oe7&
z^zDMG6ha04niUt@`7`q+&9F)TU*i@s%fkB?vK~GWo-d<yCG}NgIV}GowQyH6xSqpj
z4lMtY+!?r2ojUJ|CMWOlZ*V6T$nb<=`y^tNi*`US@toMQDef`ZXszwe9#eVZm2vHr
zr0+AUyQZ?r>ubyEpudM2ynN4SN!eMh_D&nsw||nu)B)#`wK!wlGJ2Y-&~Bjb*k^{|
ze#z(e?C4KdD4NpAL<WQ!)y?M9qj)4(eBn%>ug;ZHHWcs_$KP#uzkv8KS9oAaoB2Hk
z`@jxeANWC^7W4}%P{-_9iczllaBAsQGOYmzO<LAR4$ZRpNP;NY9}9Y_WL|5bH4S;E
zFG7i|)|t}gv<6!`mW_pvg=zlYHHn`Gz3$UrC}@tH^ABW@JBLvrGQuJ9A6g*Zx+0rZ
z3VF+i1N3<<&I>Bbd;II7p1NfWNDHd;el{Dm0+z0$YxkytIru=I!uUXPrcF1FcF`v$
zF<z-CSnQQVf7kCt0i`eF|8l~QKyP1UT@m$!-%W}dH_>)FvMW*eB)lN>v#C%J!MxJ1
zS^)Hml;4xd_okG+sh9as9EDET8J>{4<$CJqxao=7<#}+>!yM?;{6Cvb{)7pR#;lRZ
zzHGR{9OC&<=I62<otblI{tw|Yb6-}L28XU5t`WFag?{(^&&gU{$@*v|+LP4!w3{95
z+f?<o2#DTX88mKZF}Mn8X!-cZypK&?>>%N}(mhwW2IxoLsndux+E6Y-;uXmC#nVTg
zXeHmSXDS}=9tWbE*3P`Qc581{vU_Q!p~9O>1Mne{!A8Obw7@EVOwmV~P3CNBu#ZU3
z$ZKFu@n+@Ny2K?%Gj={tVJA0=XG-opXI<|7k9wdlA^VQmMDA#X)|?#G3^tpyYCVY0
zEN{w7awk8_VILC({edTYG#~6RxIj^No}L?>z%F;8nJUX=lL@jFeOqoe2J{v%$eqij
z&jKq;XC2~Q`@FALcQtoJ59!fKByJ^S6u$tyfH|$~0?a3*2Qep7dSc?Vzpp8wP~^LW
z{ZJ-y-;DJ!=yOEXf&#G*9$SdbmOuor4_u)>M7wbZ85g3lenIp>XU;3V_90)Q>Qu9G
zB+M6*vdMV6&%yZbgTTwF1P`tN+!7VoU(lb*poP;Rn%S?#gv+PkEFHy1>3fJ-cb3~m
zTu1~3(5I)iiDPVnFeZ<3y!SXoH+^8?W=wmqbq$C}H{KKd74-V_dJ(ENbC1TGEh;_d
zPY3Mm$GmUfrawG=I>SJsgnR=1zteus^|^M-gHxG&t!c&FBw0A$(58ys7ifC%tZHl{
z=(kG6)zv(IABz;$Hu~D9%+}xO&A`KA)gVygk;FxXGw&5-=MO~TlNDZ3r=aIIZ{|Ne
zHjDevs$LtZ(|6JwaW%4FpV2IPPGau95N+f+ztU=@4r~AoWa+Xe|E(a;4oJN03FtQ*
zd`C(78-6aW;(EueK9zy(t(!@ET&|x91-;fYF-4%KPf#2?9cN6!UtGiN<d^+78~#HW
z$vpl_B=5OOr^fY1pg-9=M7^N$US&Sew(SV1X@5~V6yQ($wav4;=g~DRJ2U8Ua+Y+C
zOMYB2am}XNRYm&xk<=njuU-u64ZzXo+oj}yzT<6$SLMt#J|6{5uRF4gf^1EnkBldo
OFuNs&3d|32tp5jf1uQoJ

diff --git a/ee/audit/test_out.txt b/ee/audit/test_out.txt
deleted file mode 100644
index c406046d4a5253828013b4c4039604002ee0ab47..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 4768
zcmV;R5?}2VBilxmY@Fntd^C?J<zg#r<K|+cxP!pBGyDg7jIj))fnVBcMAa)IOVjX1
z;hCmxIX>gD&7ha1S-mA9u*_CiC$ZIUNvC>BF?5ecUqY<sg9!lHtQ?SzQ(1t1N`(3K
zxWZc+e72u8KB-Cgq0*^F(I#9~Kx>q~sGm}AsCbB$W!L<=1r7d#OVK*0g2+-%+(ss=
zp)XtS1dw~azCh9NF4~5gc`Rkta7w#vy>9K48%VFkK{j6kW^R3#J7hie$|0|bd6z(j
z=jGARrk#?1`V$Js=<rd?mj}w06#LZv&(n_p&&B`_Di2HvUy-u8;6W?Q$IOoSJagyM
z-Q9!Hdqr_g#xk@EbWaT~%IdRfdGDn{p5H1nRZjuHB0E5i_m`Q1{z572)XGOh0`xS=
zK-Jw+K#1*h`@I4yY3G}Y%=ll9lV3K__M6Szefms5?FF#*R@eCDa-S$*U|oUV->R!L
zuWc)dEvK@rIhsuR+_<TP!kcobaA^vh6!rXop;|>cW1VTsyDFl5s8w%0F0VF8JOMca
zgAlBAdl?>VJTEI%Y(~w(9sSDBtL?-!_sK=P#(FV`HQ0P^O~q=k*DdWJ{Vh7i=4{A+
zu5xC&UH-BhZ*W0;()+7m7_uIphEh^E1*EwJ`-NpNTF>rnynHkXERFCH4<)2ZMtMQx
z^pH83LYcJI{p2Qoe$m|x6(iM^3B)0oZpVvb`Rmo>ZqEo`|3I$Uz{mkc#^Q|rEPE{!
zG2TjqQFJ_W{@Oc!=S_u{0tvZPm=zp>Z4nl@O{Srx!=vc$o(u14;(T>`@ckmMSMK+;
zPbHw*ARWXPAdHXKI_j1svb=_TEb>~K7*VN<-8c~#nljp~ruh@fEl#qMEqBQrP&4Lg
zP<=%#i42V_@T@;CuZM|8esu`pJ$f4U@LEc#5!Zedj}Yf2As?Tv-&ye0i6xViL8Po-
zfUz2v1>ZVS6WX${Fi&{*ilSZa%@mu+$Q&Ym06GK^h%*N_^aS3NbaB>{2PxeGUPq8<
z0ot`RFy%TOPpl=1tU{U>Ej>{eOXz_v6CeB}w5hY=>KyU|g_#@bf*qASF1X}y*a@tl
zsSou=@x~XrQ);KxF&h=Bjz6<Jq_S_UA{^{2?gl1~ODdkz{sGqiOS&RuvEWa#+jH%o
zgS%z|%EME^X~-aJM2o$g5L~r_leW%hr$5X`lKvi`rW9Z1xT|B8Bap|CY#T2%8_6X)
z==|${K|63B<jyS!KIvQ87e1rh=UpFFmoMMrX^wY|pP%l59e$Duq_A*xVZlwR&a5rw
zLT+M%vN!}a3<Y*4V@}WHM5g$@pjc$Jw?rY--gQoEo3~03_gQx5Yz%tTP9uoBW<l37
zD<^6mtUnlub$sr2mnl9i+arwP&Lk#14J?_+GE!{Z@+zsb2&}$30uA4?hWR%(LkIzB
zmhG;Os4psb%t&2P5z82P{}Wd9wtMXujI(4>gXkWJfT|l)M~Y?0S|qMCmv@ZUn~Ued
zseccO)l%ovt&V-|OpWXmG@ocugwshEEJMdfGaiAKKe;CD+TT><pD1G2Uf8wr8@5Ux
z*U%K<&r=I<u~7^|A*xzA6LK^qz8zU)j5MKyI)rY-`BwEJ1R1txr~`M!eI8&5;=dJ~
z(%WwT>CKzyeEj?io+(!=6GtB!#h*MO-Bo&ffljZYwHuQ-t#lXRR^w8*H2+RZVhb6M
zQGmjKTm4#$U2c4O!Pcq1QRf-K)up>>tOnQChFeZslH}O)bHaXqJ^}1{@yl1<Pr)@l
zlZ5VuL<T$Ro=u$e4UJ`GgR=d4K5NVdEFts0PK!1AUP{{!y@Wl9f#hL^w+QReY!w|k
zGf|y|Wpt-Y^TGc%6GfCmLnE;+a#<}%;uNZ{+`$3LgH|0>xF4icXNMiQ8bQCk3IZYz
z2Gq-;eQA*dgj$7Rf$p{Dve0FD9@;W$t$L6+e)H*=wjspx>h*Ts#M}&_a(254Ql=53
z(X`@!<Zh(6{IZhwx**;{P`Opy1@Rl%6q11tTEI`~q{YU(6`J5X<}!TeR2LfYgIEUB
z+Iqn^pua%83=HKO75KG&GhcqmSp3r#p|5u+<21B&n4fVoL!_SRWzI6j?=qCuPj9Pq
z2PR&V+Etr2hv0^<QDupouV%7F)J{IFP);+P8{tGBzj<;lE2qbg>c-?~K<QzCXtjxr
zk=n6iW3221o{g`_eXd5$p2!_!$)sM*b@|HofDFKaz=WV)b&HuAiLdOY!hfO3jVN20
z8z;OJww04I-o_}AD5SvaHsQZNB~yFN7RWrT&y$I7f=6T(LCwVZcglo*7C~=|C|Uj>
z%Z2seNyU78_LE0I4JhCs1HnSxJ5q=?Sn=37L*%bRAr98kqaJjU#0Qs5^-_8a#{bJ&
zV@DJTHPt`MhnE~3Nm93t7JxWKKNzDa+oiy?HeIOBcCS~|Sl7R=PY25V0*tR?C3sRw
z_KAP0HcE2<d4|bk#*<<DJbzByG$U_vLj0#(W0D<L_GfvR9|JaV1P8Ygr5mf7soiQH
z9}!p{oEI*9UMj_|%8RrYuViAA@koczw1Iy(50<onbH~GkMNK!(iC6f7AS<=$9_&_W
zq9>>0c{mzMzRWH!);w0M^;x`igxex|$?;}6Ty>;1KU7R{21=t`c2;l?t?$Elj*@%8
zL`6FC_0n$JZX^<*rwVN$;i~Lq`aVDR6=%cf$Ec9aiQWtbwMckuTRj4<CJ!!br+Bh}
zCds<M$r(Y)!M!qh--)q(h-9%S-j#LF-=tb=><6&B&AAgspp;8uU!{RhItFHlJkX%@
zzGK1Qo~T0y3s$BAZb~>=uyHL87ho1dLI3XlsKDk<jvyNmh0mN$Yq-dRz;w))szkq#
zOrIi`WhGdLqKjK64Mo39z<zrlLWOEq);y}kqQo&)$Ry^=HVx8yqE|@g9|D2$=u>DM
z-_Vc}wh7H4Ex#&mK%p9Nv6NuN>34=6Sr7I48<7LQ7oZ91dhq9Tb4HcgcDYX|{{$AO
zdO^v=yG5h7c;i4Q<0wdAC3OBb2zS1ISw>c%8`zfk&J<xnuaJd5ObWy>7Oz?x7|=g_
zEr_2|<B)v@9rcq6&#k@1z~i=mn<1@<aP>|fqB!$GGL!lMnq_vHW*<&AYU5Fxc3cgv
z5MWyIPoIho`h1hB_G;oHMn0I<O$C09r+G9FB-3@9O7w&s0M}YSCGOQm3Dqq|bd_nt
zHlT%G-4WCo%=YOrYP_(p+9QU8Bbor^8sqhJ)SDB{NiH!U;EQyb+{-a75G3v8h3(Aq
z<d!dvjID%J!cDg-ue9<eSPTh8uvWMhQd=_MRJX`EVm|EqIXoL-jun|Po)Q2@>Epge
zdd~&xwp@ME2#2hrZaRZ2hg;9w-GxFQ$kTDR)|LlY0k-+wA&^CCH2AbS6?yZa1zX)B
z8Twz^=j)F$?xtuF&}ViMJo5B)Nbz8hY@2#!qer?B;Kp7>^CAuQlQTtlQJ=^ee~HoJ
z_VbnxBYd4(>L3c1h>7Y|HfvSKfPB<`{U_7-`9v#e-X&%B0}o5Nr$y#46Df^cxI|S?
zpE8y-QYdn*7sRiZEpU-;&9c0Aii1bcIRBs>d&)^4kWq+rfHhZe_O09zERd(B<QLJI
z`gc9^{BCw9D?eTcy{yG51%At+mH|K)#JN~Rx6{^hAvr7_xwi*KjvXMoJjzJ}imQ!0
z4^OdKo(iEf!m8MIlpL?#sQrOAEGQr+FSg@Lv#M4K$Pl9Ib4loDDGlhFnj6y;Wb0#T
zsTu7HrWKMSC~pM+ud)N}M;b+QC8&C5%61mY*()=oIH;?ojY9_@?C^rS3F!b+)AT7M
z>SmLY=Ig-56~z`e53tVs^SCThPRY&02;I)92oYJrj^y%4T3(Ep4Z63vDsyCr>XGhA
z{oj<c(qsDghXK{|1{Nxj`*~k2xy#cvJ>eP(;|_pZ^(;ZGCbB4W&wcSJ-6nqE&jaMi
zoC~C))dgN}ZJ>ru7|ZLbul_)UKq!D}-7GM!z5Bn!zx*#LA=d9q)2N%W_G)oY6QvLw
z{pS#FLhsO0ED0hrM?rit@*AqN$N&T0uwN`@NF^lO+#PvfPqn5F<5;aTuYOFEz^I2r
z?}hZnG;M0EpW~W&k+BvH7iTbT{F-J#VBE!L|HuI=a`=$$;^VxuEU_;)5yQjzGJ5B`
z<|5m1PLa=aytt}tsXe)W%Mo9C$2pMA>o+3m-U_xvx%jL1ev~W*^Vo?|3D6(xyG-{*
zvYbhD9uK}F$K?vrXV0oEA?!4$m;(kTfns}{*Q`hLJn7a~(n3GV1y}S&f8vyIuW?=%
z4m7C-acAL%uV5d~JT4W|#J|mC!kcM+gUTCYN6S*1xD6DH1IlhKhlieU55yDPgqL%>
z;jG%1O)F_MZXW~>euFC}hnqLb{F6zEq?XT^&577;g}@tp-@Os@hu;sKrcyT<B;>n<
zAiaG<4L>>T)p$ikoCeIiw#+M^x8J2sKO{MqK}sTTSy_A5FEDctEbm^lfW{s`PC78b
z$`69xtU%sPNvN0w=A{Zu)*PsSL=OCK_+9E4CBBV*^a{O&RRZ#x2JB;<>;_-|q6;%}
z@Wa6)@J#rcvk$t(D_TcqB{_@(oo%0ub>4?Op-4*w+Di{s#dO|PbZn*cF8-<Dl691O
z`DV!+gw#H{o9cMXr*C=STAK~4szYP`Dbn2^v7Z~O4#HE8K8gV!61VO2V-s#yL;Y_+
zMCuLe70leC=S|7u8IFMFe0aG=3GZ%vmhTieIh<>Fgvkw5z&J=0_DPDF`YDSQGN}9F
zyUc*YQ(G#?enuS;q-e2{Sp*?m1<`g@UN_4S|C-RshagYW(!@dN#7ElDt{6<SqB!uT
zpqwAh4|*Awmn*a{eRLuW;^|tlL#xhv?ydjr{I&2s<UzaZ<A(i!762ds7Wp+qS^Z`z
z=sgzeR}YFghYVeK%73x&^{pGn33+EJ?o|rDMB<L!h=YitZQb5_8E6d8cHg%*^CaXD
zXX4s1YkLs*0HaE!Q~JtJXI^r=HR2v7qAz_8UI>X4l5HS7!4_z9r^BPL;v1sSuqNqV
z!nsZ-v+Z^1iJ$u|c0lvklZY>ufEtN}Ygrl?x7Sh%2~tpiM0SHvz~$_@1XZ2#)g3B(
zS1Ud|zmMx+|03UL$^=64G3r=+^#OMz<g9MT5~`-B>V`-6>{)O3L1WV-J@74Y3PW85
zKgIh#xcZF!OhPo2MnZeRnyhbsmknyy_1wmr$|pWn2$_fW85y&x$sIMOf+s2SceBOt
zZL@YKdZ9<KBf~!1kD;FJtdLsOhn4OH`tftkWxQ5qJrxPuw-F?=r!*a9V65Jij04_P
zt$mdIj$8D9g6ob`DKYBe0#okB>rmFnO4%vGmSFOwVK2ph1DY%dUbOGD<58_{<6Ixe
zD}1Za%b%hbF6;U*M$~vb!{~ZQ>f!z!QrpfGi#$oOs5Vl6%|az$WN8zz^HpuC7b+2M
z@vhFU9tqFfv}x05edEq$;l!^f0JMs-hPLq&YE;ko>i&UhDegC(gH1?|ikw=DW!`Zu
zo42a=dG-hWDCUb<r98neqAZ&A^xWQCp#KT|nQheSH7r00!%_meZ!wFwv+7ED?J+i;
z*Sv=eqFDm+kMamQ#R#Ls<N;7a@20ko%Xy<w&QmjCDA4UWWzrRRM8J2WHX|9bl^G2S
zk+^Y#WosPcdrU#a_uzYnj=G<4|3oevP6;G$5p^dZxDr7zytwm9mrdcK)k7vWRF`Ex
zI+D_v?n+VDhFOR)<2fjUO8uewfCtpjV`-&dSaVkG^kT#L<2qANhPX@1i=x&SXxDRz
zN?Ra$E#<~F{sJK5NF@<u@D6W)zb-*1!!Q5m`0nrQOTzGRf>+72K@8o(XdVc|@9>oG
z$p{B{`!{!pGecb)j^JFW>m_XSpNW%WQ8FwNXn+5VqRJi#)mr`Yn91!zn0c>(JYfyA
zEyzXmA8JZ+GJ`-p{Qo2<xizRb#U7mZzdIdH$(C2(bL@lf89f<Pnfzx+M{kukc>(5O
zspkP&Za0KO*Qe|f!m44on>E!rV70Zq4=uWok%jRPrRW|u)LK4?_H_2)I_s|kWnlcA
zw30cki8=M8b(SaY`DN({dCt=rIdF4GYDwnj#U$pJN$$_tC(}uA(CUj{H4M>B^VCgt
zQi;@Ipe|?8a7c=K7W>4l99=ANV?`pdqKc1qg4vg@DC7duP}G#*>me**f=EgkY8r|0
ukj%(<f9$;=^I@cmJd|=<c$63i9^}*YZc;h8gyQ;LcloEV3yBA*#QS@J9Y*K?

diff --git a/ee/utils_ee.go b/ee/utils_ee.go
index cf2007a5c6b..d2788dda918 100644
--- a/ee/utils_ee.go
+++ b/ee/utils_ee.go
@@ -37,7 +37,7 @@ func GetEEFeaturesList() []string {
 	} else {
 		ee = append(ee, "backup_restore")
 	}
-	if worker.Config.AuditEnabled {
+	if x.WorkerConfig.Audit != "" {
 		ee = append(ee, "audit")
 	}
 	return ee
diff --git a/go.mod b/go.mod
index 63b6afaa57e..2cee226778a 100644
--- a/go.mod
+++ b/go.mod
@@ -5,7 +5,6 @@ go 1.12
 // replace github.com/dgraph-io/badger/v2 => /home/mrjn/go/src/github.com/dgraph-io/badger
 // replace github.com/dgraph-io/ristretto => /home/mrjn/go/src/github.com/dgraph-io/ristretto
 
-replace gopkg.in/natefinch/lumberjack.v2 => /home/amanbansal/go/src/gopkg.in/natefinch/lumberjack
 require (
 	contrib.go.opencensus.io/exporter/jaeger v0.1.0
 	contrib.go.opencensus.io/exporter/prometheus v0.1.0
@@ -70,7 +69,6 @@ require (
 	golang.org/x/tools v0.0.0-20210106214847-113979e3529a
 	google.golang.org/grpc v1.23.0
 	gopkg.in/DataDog/dd-trace-go.v1 v1.13.1 // indirect
-	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 	gopkg.in/square/go-jose.v2 v2.3.1
 	gopkg.in/yaml.v2 v2.2.4
 )
diff --git a/go.sum b/go.sum
index d300de1abbf..7be5b12ca09 100644
--- a/go.sum
+++ b/go.sum
@@ -805,8 +805,6 @@ gopkg.in/ini.v1 v1.42.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k=
 gopkg.in/ini.v1 v1.51.0 h1:AQvPpx3LzTDM0AjnIRlVFwFFGC+npRopjZxLJj6gdno=
 gopkg.in/ini.v1 v1.51.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k=
 gopkg.in/mgo.v2 v2.0.0-20180705113604-9856a29383ce/go.mod h1:yeKp02qBN3iKW1OzL3MGk2IdtZzaj7SFntXj72NppTA=
-gopkg.in/natefinch/lumberjack.v2 v2.0.0 h1:1Lc07Kr7qY4U2YPouBjpCLxpiyxIVoxqXgkXLknAOE8=
-gopkg.in/natefinch/lumberjack.v2 v2.0.0/go.mod h1:l0ndWWf7gzL7RNwBG7wST/UCcT4T24xpD6X8LsfU/+k=
 gopkg.in/resty.v1 v1.12.0/go.mod h1:mDo4pnntr5jdWRML875a/NmxYqAlA73dVijT2AXvQQo=
 gopkg.in/square/go-jose.v2 v2.3.1 h1:SK5KegNXmKmqE342YYN2qPHEnUYeoMiXXl1poUlI+o4=
 gopkg.in/square/go-jose.v2 v2.3.1/go.mod h1:M9dMgbHiYLoDGQrXy7OpJDJWiKiU//h+vD76mk0e1AI=
diff --git a/worker/config.go b/worker/config.go
index 490aca142ed..76506305319 100644
--- a/worker/config.go
+++ b/worker/config.go
@@ -17,6 +17,7 @@
 package worker
 
 import (
+	"github.com/golang/glog"
 	"path/filepath"
 	"time"
 
@@ -71,8 +72,7 @@ type Options struct {
 	// CacheMb is the total memory allocated between all the caches.
 	CacheMb int64
 
-	AuditEnabled bool
-	AuditDir string
+	Audit string
 }
 
 // Config holds an instance of the server options..
@@ -97,12 +97,19 @@ func (opt *Options) validate() {
 	x.Check(err)
 	td, err := filepath.Abs(x.WorkerConfig.TmpDir)
 	x.Check(err)
-	ad, err := filepath.Abs(opt.AuditDir)
-	x.Check(err)
 	x.AssertTruef(pd != wd, "Posting and WAL directory cannot be the same ('%s').", opt.PostingDir)
 	x.AssertTruef(pd != td, "Posting and Tmp directory cannot be the same ('%s').", opt.PostingDir)
 	x.AssertTruef(wd != td, "WAL and Tmp directory cannot be the same ('%s').", opt.WALDir)
-	x.AssertTruef(ad != pd, "Posting and Audit Directory cannot be the same ('%s').", opt.AuditDir)
-	x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", opt.AuditDir)
-	x.AssertTruef(ad != td, "Tmp and Audit directory cannot be the same ('%s').", opt.AuditDir)
+	if opt.Audit !="" {
+		dir := x.GetFlagString(opt.Audit, "dir")
+		if dir == "" {
+			glog.Fatal("audit flag is provided but dir is not specified")
+		}
+		ad, err := filepath.Abs(dir)
+		x.Check(err)
+		x.AssertTruef(ad != pd, "Posting and Audit Directory cannot be the same ('%s').", dir)
+		x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", dir)
+		x.AssertTruef(ad != td, "Tmp and Audit directory cannot be the same ('%s').", dir)
+
+	}
 }
diff --git a/x/config.go b/x/config.go
index 607848ba078..e4448afc7c7 100644
--- a/x/config.go
+++ b/x/config.go
@@ -112,6 +112,9 @@ type WorkerOptions struct {
 	LogRequest int32
 	// If true, we should call msync or fsync after every write to survive hard reboots.
 	HardSync bool
+
+	// Audit contains the audit flags that enables the audit.
+	Audit string
 }
 
 // WorkerConfig stores the global instance of the worker package's options.
diff --git a/x/flags.go b/x/flags.go
index 139cd168803..38d5aa8a950 100644
--- a/x/flags.go
+++ b/x/flags.go
@@ -112,3 +112,7 @@ func GetFlagUint32(opt, key string) uint32 {
 	Checkf(err, "Unable to parse %s as uint32 for key: %s. Options: %s\n", val, key, opt)
 	return uint32(u)
 }
+func GetFlagString(opt, key string) string {
+	val := GetFlag(opt, key)
+	return val
+}
diff --git a/x/log_writer.go b/x/log_writer.go
index a3dba65170e..d0f9a37f454 100644
--- a/x/log_writer.go
+++ b/x/log_writer.go
@@ -20,38 +20,68 @@ import (
 	"compress/gzip"
 	"encoding/binary"
 	"fmt"
-	"github.com/dgraph-io/badger/v2/y"
+	"github.com/dgraph-io/badger/v3/y"
 	"io"
+	"io/ioutil"
+	"math/rand"
 	"os"
+	"path/filepath"
+	"strings"
 	"sync"
+	"time"
+)
+
+const (
+	backupTimeFormat = "2021-01-19T15-04-05.000"
 )
 
 var _ io.WriteCloser = (*LogWriter)(nil)
 
 type LogWriter struct {
-	FileName      string
+	FilePath      string
 	MaxSize       int64
 	MaxAge        int // number of days
 	Compress      bool
 	EncryptionKey []byte
-	size          int64
-	file          io.Writer
-	mu            sync.Mutex
 
-	millCh    chan bool
-	startMill sync.Once
+	baseIv          [12]byte
+	mu              sync.Mutex
+	size            int64
+	file            *os.File
+	mch             chan bool
+	startDirManager sync.Once
 }
 
 func (l *LogWriter) Write(p []byte) (int, error) {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
+	var err error
 	if l.file == nil {
+		l.manageLogDir()
 		// open file
+		err = l.open()
+		if err != nil {
+			return 0, fmt.Errorf("not able to create new file %v", err)
+		}
 	}
 
-	if l.size+int64(len(p)) >= l.MaxSize {
-		// rotate
+	if l.size+int64(len(p)) >= l.MaxSize*1024*1024 {
+		err := l.rotate()
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	// if encryption is enabled store the data in encyrpted way
+	if l.EncryptionKey != nil {
+		bytes, err := encrypt(l.EncryptionKey, l.baseIv, p)
+		if err != nil {
+			return 0, err
+		}
+		n, err := l.file.Write(bytes)
+		l.size = l.size + int64(n)
+		return n, err
 	}
 
 	n, err := l.file.Write(p)
@@ -59,7 +89,6 @@ func (l *LogWriter) Write(p []byte) (int, error) {
 	return n, err
 }
 
-
 func (l *LogWriter) Close() error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
@@ -70,33 +99,119 @@ func (l *LogWriter) Close() error {
 	return nil
 }
 
-func encrypt(key, src []byte) ([]byte, error) {
-	iv, err := y.GenerateIV()
+func (l *LogWriter) open() error {
+	err := os.MkdirAll(filepath.Dir(l.FilePath), 0755)
 	if err != nil {
-		return nil, err
+		return err
+	}
+
+	size := func() int64 {
+		info, err := os.Stat(l.FilePath)
+		if err != nil {
+			return 0
+		}
+		return info.Size()
+	}
+
+	openNew := func() error {
+		f, err := os.OpenFile(l.FilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
+		if err != nil {
+			return err
+		}
+		if l.EncryptionKey != nil {
+			rand.Read(l.baseIv[:])
+			if _, err = f.Write(l.baseIv[:]); err != nil {
+				return err
+			}
+		}
+		l.file = f
+		l.size = size()
+		return nil
+	}
+
+	info, err := os.Stat(l.FilePath)
+	if os.IsNotExist(err) {
+		return openNew()
+	}
+	if err != nil {
+		return err
+	}
+
+	// encryption is enabled and file is corrupted as not able to read the IV
+	if l.EncryptionKey != nil && info.Size() < 12 {
+		return openNew()
+	}
+
+	f, err := os.OpenFile(l.FilePath, os.O_APPEND|os.O_RDWR, os.ModePerm)
+	if err != nil {
+		return openNew()
+	}
+
+	// If not able to read the baseIv, then this file might be corrupted.
+	// open the new file in that case
+	_, err = f.ReadAt(l.baseIv[:], 0)
+	if err != nil {
+		return openNew()
+	}
+	l.file = f
+	l.size = size()
+	return nil
+}
+
+func (l *LogWriter) rotate() error {
+	var err error
+	// file not open
+	if l.file == nil {
+		return l.open()
 	}
+
+	if err = l.file.Close(); err != nil {
+		return err
+	}
+
+	_, err = os.Stat(l.FilePath)
+	if err == nil {
+		// move the existing file
+		newname := backupName(l.FilePath)
+		if err := os.Rename(l.FilePath, newname); err != nil {
+			return fmt.Errorf("can't rename log file: %s", err)
+		}
+	}
+
+	err = l.open()
+	l.manageLogDir()
+	return err
+}
+
+func backupName(name string) string {
+	dir := filepath.Dir(name)
+	prefix, ext := prefixAndExt(name)
+	timestamp := time.Now().Format(backupTimeFormat)
+	return filepath.Join(dir, fmt.Sprintf("%s-%s%s", prefix, timestamp, ext))
+}
+
+func encrypt(key []byte, baseIv [12]byte, src []byte) ([]byte, error) {
+	iv := make([]byte, 16)
+	copy(iv, baseIv[:])
+	binary.BigEndian.PutUint32(iv[12:], uint32(len(src)))
 	allocate, err := y.XORBlockAllocate(src, key, iv)
 	if err != nil {
 		return nil, err
 	}
-
-	allocate = append(allocate, iv...)
-	var lenCrcBuf [4]byte
-	binary.BigEndian.PutUint32(lenCrcBuf[0:4], uint32(len(allocate)))
-	allocate = append(lenCrcBuf[:], allocate...)
+	allocate = append(iv[12:], allocate...)
 	return allocate, nil
 }
 
 func compress(src string) error {
 	f, err := os.Open(src)
 	if err != nil {
-		return fmt.Errorf("failed to open log file: %v", err)
+		return err
 	}
 
 	defer f.Close()
-	gzf, err := os.OpenFile(src+".gz", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
+	gzf, err := os.OpenFile(src+".gz", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, os.ModePerm)
 	if err != nil {
-		return fmt.Errorf("failed to open compressed log file: %v", err)
+		return err
 	}
 
 	defer gzf.Close()
@@ -104,17 +219,110 @@ func compress(src string) error {
 	defer gz.Close()
 	if _, err := io.Copy(gz, f); err != nil {
 		os.Remove(src + ".gz")
-		err = fmt.Errorf("failed to compress log file: %v", err)
 		return err
 	}
-
+	// close the descriptors because we need to delete the file
+	if err := f.Close(); err != nil {
+		return err
+	}
 	if err := os.Remove(src); err != nil {
 		return err
 	}
 	return nil
 }
 
-func manageOldLogs() {
-	// delete old log files.
-	// compress log files
-}
\ No newline at end of file
+func (l *LogWriter) manageLogDir() {
+	l.startDirManager.Do(func() {
+		l.mch = make(chan bool, 1)
+		go func() {
+			for range l.mch {
+				l.manageOldLogs()
+			}
+		}()
+	})
+
+	select {
+	case l.mch <- true:
+	default:
+	}
+}
+
+// this should be called in a serial order
+func (l *LogWriter) manageOldLogs() {
+	toRemove, toKeep, err := processOldLogFiles(l.FilePath, l.MaxSize)
+	if err != nil {
+		return
+	}
+
+	for _, f := range toRemove {
+		errRemove := os.Remove(filepath.Join(filepath.Dir(l.FilePath), f))
+		if err == nil && errRemove != nil {
+			err = errRemove
+		}
+	}
+
+	// if compression enabled do compress
+	if l.Compress {
+		for _, f := range toKeep {
+			// already compressed no need
+			if strings.HasSuffix(f, ".gz") {
+				continue
+			}
+			fn := filepath.Join(filepath.Dir(l.FilePath), f)
+			errCompress := compress(fn)
+			if err == nil && errCompress != nil {
+				err = errCompress
+			}
+		}
+	}
+
+	if err != nil {
+		fmt.Printf("error while managing old log files %+v\n", err)
+	}
+}
+
+func prefixAndExt(file string) (prefix, ext string) {
+	filename := filepath.Base(file)
+	ext = filepath.Ext(filename)
+	prefix = filename[:len(filename)-len(ext)]
+	return prefix, ext
+}
+
+func processOldLogFiles(fp string, maxAge int64) ([]string, []string, error) {
+	dir := filepath.Dir(fp)
+	files, err := ioutil.ReadDir(dir)
+	if err != nil {
+		return nil, nil, fmt.Errorf("can't read log file directory: %s", err)
+	}
+
+	defPrefix, defExt := prefixAndExt(fp)
+	// check only for old files. Those files have - before the time
+	defPrefix = defPrefix + "-"
+	toRemove := make([]string, 0)
+	toKeep := make([]string, 0)
+
+	diff := time.Duration(int64(24*time.Hour) * int64(maxAge))
+	cutoff := time.Now().Add(-1 * diff)
+
+	for _, f := range files {
+		if f.IsDir() ||
+			!strings.HasPrefix(f.Name(), defPrefix) ||
+			!strings.HasSuffix(f.Name(), defExt) ||
+			!strings.HasSuffix(f.Name(), defExt+".gz") {
+			continue
+		}
+
+		p, e := prefixAndExt(fp)
+		ts, err := time.Parse(backupTimeFormat, f.Name()[len(p):len(f.Name())-len(e)])
+		if err != nil {
+			continue
+		}
+		if ts.Before(cutoff) {
+			toRemove = append(toRemove, f.Name())
+		} else {
+			toKeep = append(toKeep, f.Name())
+		}
+	}
+
+	return toRemove, toKeep, nil
+}
diff --git a/x/log_writer_test.go b/x/log_writer_test.go
new file mode 100644
index 00000000000..c47ae42999a
--- /dev/null
+++ b/x/log_writer_test.go
@@ -0,0 +1,157 @@
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package x
+
+import (
+	"bufio"
+	"bytes"
+	"compress/gzip"
+	"crypto/aes"
+	"crypto/cipher"
+	"encoding/binary"
+	"github.com/stretchr/testify/require"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestLogWriter(t *testing.T) {
+	path, _ := filepath.Abs("./log_test/audit.log")
+	defer os.RemoveAll(filepath.Dir(path))
+	lw := &LogWriter{
+		FilePath: path,
+		MaxSize:  1,
+		MaxAge:   1,
+		Compress: false,
+	}
+
+	writeToLogWriterAndVerify(t, lw, path)
+}
+
+func TestLogWriterWithCompression(t *testing.T) {
+	path, _ := filepath.Abs("./log_test/audit.log")
+	defer os.RemoveAll(filepath.Dir(path))
+	lw := &LogWriter{
+		FilePath: path,
+		MaxSize:  1,
+		MaxAge:   1,
+		Compress: true,
+	}
+
+	writeToLogWriterAndVerify(t, lw, path)
+}
+
+func TestLogWriterWithEncryption(t *testing.T) {
+	path, _ := filepath.Abs("./log_test/audit.log")
+	defer os.RemoveAll(filepath.Dir(path))
+	lw := &LogWriter{
+		FilePath:      path,
+		MaxSize:       1,
+		MaxAge:        1,
+		Compress:      false,
+		EncryptionKey: []byte("1234567890123456"),
+	}
+
+	msg := []byte("abcd")
+	msg = bytes.Repeat(msg, 256)
+	msg[1023] = '\n'
+	for i := 0; i < 100; i++ {
+		n, err := lw.Write(msg)
+		require.Nil(t, err)
+		require.Equal(t, n, len(msg)+4, "write length is not equal")
+	}
+
+	time.Sleep(time.Second * 10)
+	require.NoError(t, lw.Close())
+	file, err := os.Open(path)
+	require.Nil(t, err)
+	defer file.Close()
+	outPath, _ := filepath.Abs("./log_test/audit_out.log")
+	outfile, err := os.OpenFile(outPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	require.Nil(t, err)
+	defer outfile.Close()
+
+	block, err := aes.NewCipher(lw.EncryptionKey)
+	stat, err := os.Stat(path)
+	require.Nil(t, err)
+	iv := make([]byte, aes.BlockSize)
+	_, err = file.ReadAt(iv, 0)
+	require.Nil(t, err)
+
+	var iterator int64 = 16
+	for {
+		content := make([]byte, binary.BigEndian.Uint32(iv[12:]))
+		_, err = file.ReadAt(content, iterator)
+		require.Nil(t, err)
+		iterator = iterator + int64(binary.BigEndian.Uint32(iv[12:]))
+		stream := cipher.NewCTR(block, iv)
+		stream.XORKeyStream(content, content)
+		//require.True(t, bytes.Equal(content, msg))
+		_, err = outfile.Write(content)
+		require.Nil(t, err)
+		if iterator >= stat.Size() {
+			break
+		}
+		_, err = file.ReadAt(iv[12:], iterator)
+		require.Nil(t, err)
+		iterator = iterator + 4
+	}
+}
+
+func writeToLogWriterAndVerify(t *testing.T, lw *LogWriter, path string) {
+	msg := []byte("abcd")
+	msg = bytes.Repeat(msg, 256)
+	msg[1023] = '\n'
+	for i := 0; i < 10; i++ {
+		go func() {
+			for i := 0; i < 1000; i++ {
+				n, err := lw.Write(msg)
+				require.Nil(t, err)
+				require.Equal(t, n, len(msg), "write length is not equal")
+			}
+		}()
+	}
+	time.Sleep(time.Second * 10)
+	require.NoError(t, lw.Close())
+	files, err := ioutil.ReadDir("./log_test")
+	require.Nil(t, err)
+
+	lineCount := 0
+	for _, f := range files {
+		file, _ := os.Open(filepath.Join(filepath.Dir(path), f.Name()))
+
+		var fileScanner *bufio.Scanner
+		if strings.HasSuffix(file.Name(), ".gz") {
+			gz, err := gzip.NewReader(file)
+			require.NoError(t, err)
+			all, err := ioutil.ReadAll(gz)
+			require.NoError(t, err)
+			fileScanner = bufio.NewScanner(bytes.NewReader(all))
+			gz.Close()
+		} else {
+			fileScanner = bufio.NewScanner(file)
+		}
+		for fileScanner.Scan() {
+			lineCount = lineCount + 1
+		}
+	}
+
+	require.Equal(t, lineCount, 10000)
+}
diff --git a/x/logger.go b/x/logger.go
index 014f92805dc..c2927ad5942 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -22,25 +22,26 @@ import (
 
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
-	"gopkg.in/natefinch/lumberjack.v2"
 )
 
 func InitLogger(dir string, filename string, key []byte) (*Logger, error) {
 	if err := os.MkdirAll(dir, 0700); err != nil {
 		return nil, err
 	}
+	if key != nil {
+		filename = filename + ".enc"
+	}
+
 	path, err := filepath.Abs(filepath.Join(dir, filename))
 	if err != nil {
 		return nil, err
 	}
 	getWriterSyncer := func() zapcore.WriteSyncer {
-		w := &lumberjack.Logger{
-			Filename: path,
+		w := &LogWriter{
+			FilePath: path,
 			MaxSize:  100,
 			MaxAge:   30,
-		}
-		if key != nil {
-			w.EncryptionKey = key
+			EncryptionKey: key,
 		}
 		return zapcore.AddSync(w)
 	}

From cdd76e83d3644d0fda828dd049485c7fabbbe45c Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Tue, 19 Jan 2021 23:05:29 +0530
Subject: [PATCH 18/31] adding test cases to check requests are getting logged
 into the audit logs

---
 systest/audit/audit_dir/.gitkeep    |   0
 systest/audit/audit_dir/aa/.gitkeep |   0
 systest/audit/audit_dir/za/.gitkeep |   0
 systest/audit/audit_test.go         | 126 ++++++++++++++++++++++++++++
 systest/audit/docker-compose.yml    |  39 +++++++++
 x/log_writer_test.go                |   2 +
 6 files changed, 167 insertions(+)
 create mode 100644 systest/audit/audit_dir/.gitkeep
 create mode 100644 systest/audit/audit_dir/aa/.gitkeep
 create mode 100644 systest/audit/audit_dir/za/.gitkeep
 create mode 100644 systest/audit/audit_test.go
 create mode 100644 systest/audit/docker-compose.yml

diff --git a/systest/audit/audit_dir/.gitkeep b/systest/audit/audit_dir/.gitkeep
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/systest/audit/audit_dir/aa/.gitkeep b/systest/audit/audit_dir/aa/.gitkeep
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/systest/audit/audit_dir/za/.gitkeep b/systest/audit/audit_dir/za/.gitkeep
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/systest/audit/audit_test.go b/systest/audit/audit_test.go
new file mode 100644
index 00000000000..f172460ac54
--- /dev/null
+++ b/systest/audit/audit_test.go
@@ -0,0 +1,126 @@
+/*
+ * Copyright 2017-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *     www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package audit
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"github.com/dgraph-io/dgraph/testutil"
+	"github.com/stretchr/testify/require"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+)
+func TestZeroAudit(t *testing.T) {
+	zeroCmd := map[string][]string{
+		"/removeNode": []string{`--location`, "--request", "GET",
+			fmt.Sprintf("%s/removeNode?id=3&group=1", testutil.SockAddrZeroHttp)},
+		"/assign": []string{"--location", "--request", "GET",
+			fmt.Sprintf("%s/assign?what=uids&num=100", testutil.SockAddrZeroHttp)},
+		"/moveTablet": []string{"--location", "--request", "GET",
+			fmt.Sprintf("%s/moveTablet?tablet=name&group=2", testutil.SockAddrZeroHttp)}}
+
+	msgs := make([]string, 0)
+	for req, c := range zeroCmd {
+		msgs = append(msgs, req)
+		cmd := exec.Command("curl", c...)
+		if out, err := cmd.CombinedOutput(); err != nil {
+			fmt.Println(string(out))
+			t.Fatal(err)
+		}
+	}
+
+	verifyLogs(t, "./audit_dir/za/dgraph_audit.log", msgs)
+}
+func TestAlphaAudit(t *testing.T) {
+	testCommand := map[string][]string{
+		"/admin": []string{"--location", "--request", "POST",
+			fmt.Sprintf("%s/admin", testutil.SockAddrHttp),
+			"--header", "Content-Type: application/json",
+			"--data-raw", `'{"query":"mutation {\n  backup(
+input: {destination: \"/Users/sankalanparajuli/work/backup\"}) {\n    response {\n      message\n      code\n    }\n  }\n}\n","variables":{}}'`},
+
+		"/graphql": []string{"--location", "--request", "POST", fmt.Sprintf("%s/graphql", testutil.SockAddrHttp),
+			"--header", "Content-Type: application/json",
+			"--data-raw", `'{"query":"query {\n  __schema {\n    __typename\n  }\n}","variables":{}}'`},
+
+		"/alter": []string{"-X", "POST", fmt.Sprintf("%s/alter", testutil.SockAddrHttp), "-d",
+			`'name: string @index(term) .
+			type Person {
+			  name
+			}'`},
+		"/query": []string{"-H", "'Content-Type: application/dql'", "-X", "POST", fmt.Sprintf("%s/query", testutil.SockAddrHttp),
+			"-d", `$'
+			{
+			 balances(func: anyofterms(name, "Alice Bob")) {
+			   uid
+			   name
+			   balance
+			 }
+			}'`},
+		"/mutate": []string{"-H", "'Content-Type: application/rdf'", "-X",
+			"POST", fmt.Sprintf("%s/mutate?startTs=4", testutil.SockAddrHttp), "-d", `$'
+			{
+			 set {
+			   <0x1> <balance> "110" .
+			   <0x1> <dgraph.type> "Balance" .
+			   <0x2> <balance> "60" .
+			   <0x2> <dgraph.type> "Balance" .
+			 }
+			}
+			'`},
+	}
+
+	msgs := make([]string, 0)
+	for req, c := range testCommand {
+		msgs = append(msgs, req)
+		cmd := exec.Command("curl", c...)
+		if out, err := cmd.CombinedOutput(); err != nil {
+			fmt.Println(string(out))
+			t.Fatal(err)
+		}
+	}
+	verifyLogs(t, "./audit_dir/aa/dgraph_audit.log", msgs)
+}
+
+func verifyLogs(t *testing.T, path string, cmds []string) {
+	abs, err := filepath.Abs(path)
+	require.Nil(t, err)
+	f, err := os.Open(abs)
+	require.Nil(t, err)
+
+	type log struct {
+		Msg string `json:"msg"`
+	}
+	logMap := make(map[string]bool)
+
+	var fileScanner *bufio.Scanner
+	fileScanner = bufio.NewScanner(f)
+	for fileScanner.Scan() {
+		bytes := fileScanner.Bytes()
+		l := new(log)
+		_ = json.Unmarshal(bytes, l)
+		logMap[l.Msg] = true
+	}
+	for _, m := range cmds {
+		if !logMap[m] {
+			t.Fatalf("audit logs not present for command %s", m)
+		}
+	}
+}
diff --git a/systest/audit/docker-compose.yml b/systest/audit/docker-compose.yml
new file mode 100644
index 00000000000..6f181dcdd4c
--- /dev/null
+++ b/systest/audit/docker-compose.yml
@@ -0,0 +1,39 @@
+version: "3.5"
+services:
+  alpha1:
+    image: dgraph/dgraph:latest
+    working_dir: /data/alpha1
+    labels:
+      cluster: test
+    ports:
+      - "8080"
+      - "9080"
+    volumes:
+      - type: bind
+        source: $GOPATH/bin
+        target: /gobin
+        read_only: true
+      - type: bind
+        source: ./audit_dir/aa
+        target: /audit_dir
+    command: /gobin/dgraph alpha --my=alpha1:7080 --zero=zero1:5080 --logtostderr
+      --audit "dir=/audit_dir" -v=2 --whitelist=10.0.0.0/8,172.16.0.0/12,192.168.0.0/16
+  zero1:
+    image: dgraph/dgraph:latest
+    working_dir: /data/zero1
+    labels:
+      cluster: test
+    ports:
+      - "5080"
+      - "6080"
+    volumes:
+      - type: bind
+        source: $GOPATH/bin
+        target: /gobin
+        read_only: true
+      - type: bind
+        source: ./audit_dir/za
+        target: /audit_dir
+    command: /gobin/dgraph zero --raft="idx=1" --my=zero1:5080 --logtostderr -v=2 --bindall
+      --audit "dir=/audit_dir"
+volumes: {}
diff --git a/x/log_writer_test.go b/x/log_writer_test.go
index c47ae42999a..64bbae5b971 100644
--- a/x/log_writer_test.go
+++ b/x/log_writer_test.go
@@ -58,6 +58,8 @@ func TestLogWriterWithCompression(t *testing.T) {
 	writeToLogWriterAndVerify(t, lw, path)
 }
 
+// if this test failed and you changed anything, please check the dgraph audit decrypt command.
+// The dgraph audit decrypt command uses the same decryption method
 func TestLogWriterWithEncryption(t *testing.T) {
 	path, _ := filepath.Abs("./log_test/audit.log")
 	defer os.RemoveAll(filepath.Dir(path))

From 87bea226f3f984a488bab7d32658536d5942352d Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Tue, 19 Jan 2021 23:21:30 +0530
Subject: [PATCH 19/31] adding interceptor ee version

---
 dgraph/cmd/alpha/run.go    |  3 +-
 ee/audit/audit.go          | 35 +++------------
 ee/audit/interceptor.go    | 68 +++--------------------------
 ee/audit/interceptor_ee.go | 88 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 100 insertions(+), 94 deletions(-)
 create mode 100644 ee/audit/interceptor_ee.go

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 4b9e6208982..7ea99295595 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -199,8 +199,7 @@ they form a Raft group and provide synchronous replication.
 		PostingListCache,PstoreBlockCache,PstoreIndexCache,WAL).`)
 
 	flag.String("audit", "",
-		`Various audit options. 
-	enable=true/false enables the audit logs (default behaviour is false).
+		`Various audit options.
 	dir=/path/to/audits to define the path where to store the audit logs.
 	compress=true/false to enabled the compression of old audit logs (default behaviour is false).
 	encrypt_file=enc/key/file enables the audit log encryption with the key path provided with the
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index 1ca78dc0b55..bde72a06de2 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -18,44 +18,19 @@
 
 package audit
 
-import (
-	"context"
-	"google.golang.org/grpc"
-	"net/http"
 
-	"github.com/spf13/viper"
-)
-
-var auditEnabled uint32
-
-type AuditEvent struct {
-
-}
-
-var auditor *auditLogger
-
-type auditLogger struct {
+func ReadAuditEncKey(conf string) ([]byte, error) {
+	return nil, nil
 }
 
-func InitAuditorIfNecessary(conf *viper.Viper, eeEnabled func() bool) {
+func InitAuditorIfNecessary(conf string, eeEnabled func() bool) {
 	return
 }
 
-func InitAuditor(dir string) {
-	return
-}
-func Close() {
-	return
-}
-
-func (a *auditLogger) Audit(event *AuditEvent) {
+func InitAuditor(dir string, key []byte) {
 	return
 }
 
-func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error) {
+func Close() {
 	return
 }
-
-func auditHttp(w *ResponseWriter, r *http.Request) {
-	return
-}
\ No newline at end of file
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 60f2563cc7a..49bf0e47940 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -1,3 +1,5 @@
+// +build oss
+
 /*
  * Copyright 2017-2021 Dgraph Labs, Inc. and Contributors
  *
@@ -17,76 +19,18 @@
 package audit
 
 import (
-	"bytes"
 	"context"
-	"io"
-	"io/ioutil"
-	"net/http"
-	"strings"
-	"sync/atomic"
-
 	"google.golang.org/grpc"
+	"net/http"
 )
 
 func AuditRequestGRPC(ctx context.Context, req interface{},
 	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-	skip := func(method string) bool {
-		skipApis := map[string]bool{
-			// raft server
-			"Heartbeat":   true,
-			"RaftMessage": true,
-			"JoinCluster": true,
-			"IsPeer":      true,
-			// zero server
-			"StreamMembership": true,
-			"UpdateMembership": true,
-			"Oracle":           true,
-			"Timestamps":       true,
-			"ShouldServe":      true,
-			// health server
-			"Check": true,
-			"Watch": true,
-		}
-		return skipApis[info.FullMethod[strings.LastIndex(info.FullMethod, "/")+1:]]
-	}
-
-	if atomic.LoadUint32(&auditEnabled) == 0 || skip(info.FullMethod) {
-		return handler(ctx, req)
-	}
-	response, err := handler(ctx, req)
-	auditGrpc(ctx, req, info, err)
-	return response, err
+	return handler(ctx, req)
 }
 
 func AuditRequestHttp(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if atomic.LoadUint32(&auditEnabled) == 0 {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		rw := NewResponseWriter(w)
-		var buf bytes.Buffer
-		tee := io.TeeReader(r.Body, &buf)
-		r.Body = ioutil.NopCloser(tee)
-		next.ServeHTTP(rw, r)
-		r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
-		auditHttp(rw, r)
+		next.ServeHTTP(w, r)
 	})
-}
-
-type ResponseWriter struct {
-	http.ResponseWriter
-	statusCode int
-}
-
-func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
-	// WriteHeader(int) is not called if our response implicitly returns 200 OK, so
-	// we default to that status code.
-	return &ResponseWriter{w, http.StatusOK}
-}
-
-func (rw *ResponseWriter) WriteHeader(code int) {
-	rw.statusCode = code
-	rw.ResponseWriter.WriteHeader(code)
-}
+}
\ No newline at end of file
diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
new file mode 100644
index 00000000000..ae43e0ec9a9
--- /dev/null
+++ b/ee/audit/interceptor_ee.go
@@ -0,0 +1,88 @@
+// +build !oss
+
+/*
+ * Copyright 2021 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
+ * may obtain a copy of the License at
+ *
+ *     https://github.com/dgraph-io/dgraph/blob/master/licenses/DCL.txt
+ */
+package audit
+
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"strings"
+	"sync/atomic"
+
+	"google.golang.org/grpc"
+)
+
+func AuditRequestGRPC(ctx context.Context, req interface{},
+	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+	skip := func(method string) bool {
+		skipApis := map[string]bool{
+			// raft server
+			"Heartbeat":   true,
+			"RaftMessage": true,
+			"JoinCluster": true,
+			"IsPeer":      true,
+			// zero server
+			"StreamMembership": true,
+			"UpdateMembership": true,
+			"Oracle":           true,
+			"Timestamps":       true,
+			"ShouldServe":      true,
+			// health server
+			"Check": true,
+			"Watch": true,
+		}
+		return skipApis[info.FullMethod[strings.LastIndex(info.FullMethod, "/")+1:]]
+	}
+
+	if atomic.LoadUint32(&auditEnabled) == 0 || skip(info.FullMethod) {
+		return handler(ctx, req)
+	}
+	response, err := handler(ctx, req)
+	auditGrpc(ctx, req, info, err)
+	return response, err
+}
+
+func AuditRequestHttp(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if atomic.LoadUint32(&auditEnabled) == 0 {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		rw := NewResponseWriter(w)
+		var buf bytes.Buffer
+		tee := io.TeeReader(r.Body, &buf)
+		r.Body = ioutil.NopCloser(tee)
+		next.ServeHTTP(rw, r)
+		r.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
+		auditHttp(rw, r)
+	})
+}
+
+type ResponseWriter struct {
+	http.ResponseWriter
+	statusCode int
+}
+
+func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
+	// WriteHeader(int) is not called if our response implicitly returns 200 OK, so
+	// we default to that status code.
+	return &ResponseWriter{w, http.StatusOK}
+}
+
+func (rw *ResponseWriter) WriteHeader(code int) {
+	rw.statusCode = code
+	rw.ResponseWriter.WriteHeader(code)
+}

From 3248e2a53c8e074fdfa5a4c72b002c0105bd1721 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Tue, 19 Jan 2021 23:42:43 +0530
Subject: [PATCH 20/31] basic refactoring and log message truncate
 functionality

---
 dgraph/cmd/alpha/run.go    |  2 +-
 dgraph/cmd/zero/run.go     |  4 +-
 ee/audit/audit_ee.go       | 82 ---------------------------------
 ee/audit/interceptor_ee.go | 92 +++++++++++++++++++++++++++++++++++++-
 ee/audit/run_ee.go         | 16 +++----
 5 files changed, 99 insertions(+), 97 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 7ea99295595..403e6b89775 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -617,7 +617,7 @@ func run() {
 
 		MutationsMode: worker.AllowMutations,
 		AuthToken:     Alpha.Conf.GetString("auth_token"),
-		Audit:      Alpha.Conf.GetString("audit"),
+		Audit:         Alpha.Conf.GetString("audit"),
 	}
 
 	x.CheckFlag(opts.Audit, "dir", "compress", "encrypt-file")
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index f15846451cb..dff4bebdfad 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -231,9 +231,9 @@ func run() {
 		}
 	}
 
-	wd, err := filepath.Abs(opts.w)
-	x.Check(err)
 	if len(opts.audit) > 0 {
+		wd, err := filepath.Abs(opts.w)
+		x.Check(err)
 		dir := x.GetFlagString(opts.audit, "dir")
 		if len(dir) == 0 {
 			glog.Fatal("audit flag is provided but dir is not specified")
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index 9ef5e0962e1..011afdc0814 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -13,15 +13,7 @@
 package audit
 
 import (
-	"context"
-	"fmt"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/metadata"
-	"google.golang.org/grpc/peer"
-	"google.golang.org/grpc/status"
 	"io/ioutil"
-	"net/http"
 	"path/filepath"
 	"sync/atomic"
 	"time"
@@ -155,77 +147,3 @@ func (a *auditLogger) Audit(event *AuditEvent) {
 		"query_param", event.QueryParams,
 		"status", event.Status)
 }
-
-func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error) {
-	clientHost := ""
-	if p, ok := peer.FromContext(ctx); ok {
-		clientHost = p.Addr.String()
-	}
-
-	userId := ""
-	if md, ok := metadata.FromIncomingContext(ctx); ok {
-		if t := md.Get("accessJwt"); len(t) > 0 {
-			userId = getUserId(t[0], false)
-		} else if t := md.Get("auth-token"); len(t) > 0 {
-			userId = getUserId(t[0], true)
-		}
-	}
-
-	cd := codes.Unknown
-	if serr, ok := status.FromError(err); ok {
-		cd = serr.Code()
-	}
-	auditor.Audit(&AuditEvent{
-		User:       userId,
-		ServerHost: x.WorkerConfig.MyAddr,
-		ClientHost: clientHost,
-		Endpoint:   info.FullMethod,
-		ReqType:    Grpc,
-		Req:        fmt.Sprintf("%+v", req),
-		Status:     cd.String(),
-	})
-}
-
-func auditHttp(w *ResponseWriter, r *http.Request) {
-	rb, err := ioutil.ReadAll(r.Body)
-	if err != nil {
-		rb = []byte(err.Error())
-	}
-
-	userId := ""
-	if token := r.Header.Get("X-Dgraph-AccessToken"); token != "" {
-		userId = getUserId(token, false)
-	} else if token := r.Header.Get("X-Dgraph-AuthToken"); token != "" {
-		userId = getUserId(token, true)
-	} else {
-		userId = getUserId("", false)
-	}
-	auditor.Audit(&AuditEvent{
-		User:        userId,
-		ServerHost:  x.WorkerConfig.MyAddr,
-		ClientHost:  r.RemoteAddr,
-		Endpoint:    r.URL.Path,
-		ReqType:     Http,
-		Req:         string(rb),
-		Status:      http.StatusText(w.statusCode),
-		QueryParams: r.URL.Query(),
-	})
-}
-
-func getUserId(token string, poorman bool) string {
-	if poorman {
-		return PoorManAuth
-	}
-	var userId string
-	var err error
-	if token == "" {
-		if x.WorkerConfig.AclEnabled {
-			userId = UnauthorisedUser
-		}
-	} else {
-		if userId, err = x.ExtractUserName(token); err != nil {
-			userId = UnknownUser
-		}
-	}
-	return userId
-}
diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
index ae43e0ec9a9..597b0243fd5 100644
--- a/ee/audit/interceptor_ee.go
+++ b/ee/audit/interceptor_ee.go
@@ -11,10 +11,15 @@
  */
 package audit
 
-
 import (
 	"bytes"
 	"context"
+	"fmt"
+	"github.com/dgraph-io/dgraph/x"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -24,6 +29,10 @@ import (
 	"google.golang.org/grpc"
 )
 
+const (
+	maxReqLength = 512 << 10 // 512 KB
+)
+
 func AuditRequestGRPC(ctx context.Context, req interface{},
 	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 	skip := func(method string) bool {
@@ -71,6 +80,80 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 	})
 }
 
+func auditGrpc(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, err error) {
+	clientHost := ""
+	if p, ok := peer.FromContext(ctx); ok {
+		clientHost = p.Addr.String()
+	}
+
+	userId := ""
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		if t := md.Get("accessJwt"); len(t) > 0 {
+			userId = getUserId(t[0], false)
+		} else if t := md.Get("auth-token"); len(t) > 0 {
+			userId = getUserId(t[0], true)
+		}
+	}
+
+	cd := codes.Unknown
+	if serr, ok := status.FromError(err); ok {
+		cd = serr.Code()
+	}
+	auditor.Audit(&AuditEvent{
+		User:       userId,
+		ServerHost: x.WorkerConfig.MyAddr,
+		ClientHost: clientHost,
+		Endpoint:   info.FullMethod,
+		ReqType:    Grpc,
+		Req:        truncate(fmt.Sprintf("%+v", req), maxReqLength),
+		Status:     cd.String(),
+	})
+}
+
+func auditHttp(w *ResponseWriter, r *http.Request) {
+	rb, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		rb = []byte(err.Error())
+	}
+
+	userId := ""
+	if token := r.Header.Get("X-Dgraph-AccessToken"); token != "" {
+		userId = getUserId(token, false)
+	} else if token := r.Header.Get("X-Dgraph-AuthToken"); token != "" {
+		userId = getUserId(token, true)
+	} else {
+		userId = getUserId("", false)
+	}
+	auditor.Audit(&AuditEvent{
+		User:        userId,
+		ServerHost:  x.WorkerConfig.MyAddr,
+		ClientHost:  r.RemoteAddr,
+		Endpoint:    r.URL.Path,
+		ReqType:     Http,
+		Req:         truncate(string(rb), maxReqLength),
+		Status:      http.StatusText(w.statusCode),
+		QueryParams: r.URL.Query(),
+	})
+}
+
+func getUserId(token string, poorman bool) string {
+	if poorman {
+		return PoorManAuth
+	}
+	var userId string
+	var err error
+	if token == "" {
+		if x.WorkerConfig.AclEnabled {
+			userId = UnauthorisedUser
+		}
+	} else {
+		if userId, err = x.ExtractUserName(token); err != nil {
+			userId = UnknownUser
+		}
+	}
+	return userId
+}
+
 type ResponseWriter struct {
 	http.ResponseWriter
 	statusCode int
@@ -86,3 +169,10 @@ func (rw *ResponseWriter) WriteHeader(code int) {
 	rw.statusCode = code
 	rw.ResponseWriter.WriteHeader(code)
 }
+
+func truncate(s string, l int) string {
+	if len(s) > l {
+		return s[:l]
+	}
+	return s
+}
diff --git a/ee/audit/run_ee.go b/ee/audit/run_ee.go
index d3411f47b9c..08fe2bf97e8 100644
--- a/ee/audit/run_ee.go
+++ b/ee/audit/run_ee.go
@@ -81,7 +81,7 @@ func run() error {
 	defer file.Close()
 
 	outfile, err := os.OpenFile(decryptCmd.Conf.GetString("out"), os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
-		0600)
+		os.ModePerm)
 	x.Check(err)
 	defer outfile.Close()
 
@@ -89,27 +89,21 @@ func run() error {
 	stat, err := os.Stat(decryptCmd.Conf.GetString("in"))
 	x.Check(err)
 	iv := make([]byte, aes.BlockSize)
-	_, err = file.ReadAt(iv, 0)
-	x.Check(err)
+	x.Check2(file.ReadAt(iv, 0))
 
 	var iterator int64 = 16
 	for {
 		content := make([]byte, binary.BigEndian.Uint32(iv[12:]))
-		_, err = file.ReadAt(content, iterator)
-		x.Check(err)
-
+		x.Check2(file.ReadAt(content, iterator))
 		iterator = iterator + int64(binary.BigEndian.Uint32(iv[12:]))
 		stream := cipher.NewCTR(block, iv)
 		stream.XORKeyStream(content, content)
-		_, err = outfile.Write(content)
-		x.Check(err)
-
+		x.Check2(outfile.Write(content))
 		// if its the end of data. finish encoding
 		if iterator >= stat.Size() {
 			break
 		}
-		_, err = file.ReadAt(iv[12:], iterator)
-		x.Check(err)
+		x.Check2(file.ReadAt(iv[12:], iterator))
 		iterator = iterator + 4
 	}
 	return nil

From 2ffcbfe53c2717026a25705cc86b834ffffefc36 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Wed, 20 Jan 2021 00:21:39 +0530
Subject: [PATCH 21/31] fixing auditing request handler

---
 dgraph/cmd/alpha/run.go    | 61 ++++++++++++++++++++------------------
 dgraph/cmd/zero/run.go     | 20 +++++++------
 ee/audit/interceptor_ee.go | 12 +++++++-
 3 files changed, 54 insertions(+), 39 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 403e6b89775..d301bcb7b72 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -433,15 +433,18 @@ func setupServer(closer *z.Closer) {
 		log.Fatal(err)
 	}
 
-	http.Handle("/query", audit.AuditRequestHttp(http.HandlerFunc(queryHandler)))
-	http.Handle("/query/", audit.AuditRequestHttp(http.HandlerFunc(queryHandler)))
-	http.Handle("/mutate", audit.AuditRequestHttp(http.HandlerFunc(mutationHandler)))
-	http.Handle("/mutate/", audit.AuditRequestHttp(http.HandlerFunc(mutationHandler)))
-	http.Handle("/commit", audit.AuditRequestHttp(http.HandlerFunc(commitHandler)))
-	http.Handle("/alter", audit.AuditRequestHttp(http.HandlerFunc(alterHandler)))
-	http.HandleFunc("/health", healthCheck)
-	http.HandleFunc("/state", stateHandler)
-	http.HandleFunc("/jemalloc", x.JemallocHandler)
+	baseMux := http.NewServeMux()
+	http.Handle("/", audit.AuditRequestHttp(baseMux))
+
+	baseMux.Handle("/query", http.HandlerFunc(queryHandler))
+	baseMux.Handle("/query/", http.HandlerFunc(queryHandler))
+	baseMux.Handle("/mutate", http.HandlerFunc(mutationHandler))
+	baseMux.Handle("/mutate/", http.HandlerFunc(mutationHandler))
+	baseMux.Handle("/commit", http.HandlerFunc(commitHandler))
+	baseMux.Handle("/alter", http.HandlerFunc(alterHandler))
+	baseMux.HandleFunc("/health", healthCheck)
+	baseMux.HandleFunc("/state", stateHandler)
+	baseMux.HandleFunc("/jemalloc", x.JemallocHandler)
 
 	// TODO: Figure out what this is for?
 	http.HandleFunc("/debug/store", storeStatsHandler)
@@ -467,8 +470,8 @@ func setupServer(closer *z.Closer) {
 	var gqlHealthStore *admin.GraphQLHealthStore
 	// Do not use := notation here because adminServer is a global variable.
 	mainServer, adminServer, gqlHealthStore = admin.NewServers(introspection, &globalEpoch, closer)
-	http.Handle("/graphql", audit.AuditRequestHttp(mainServer.HTTPHandler()))
-	http.Handle("/probe/graphql", audit.AuditRequestHttp(http.HandlerFunc(func(w http.ResponseWriter,
+	baseMux.Handle("/graphql", mainServer.HTTPHandler())
+	baseMux.HandleFunc("/probe/graphql", func(w http.ResponseWriter,
 		r *http.Request) {
 		healthStatus := gqlHealthStore.GetHealth()
 		httpStatusCode := http.StatusOK
@@ -479,20 +482,20 @@ func setupServer(closer *z.Closer) {
 		w.Header().Set("Content-Type", "application/json")
 		x.Check2(w.Write([]byte(fmt.Sprintf(`{"status":"%s","schemaUpdateCounter":%d}`,
 			healthStatus.StatusMsg, atomic.LoadUint64(&globalEpoch)))))
-	})))
-	http.Handle("/admin", audit.AuditRequestHttp(allowedMethodsHandler(allowedMethods{
+	})
+	baseMux.Handle("/admin", allowedMethodsHandler(allowedMethods{
 		http.MethodGet:     true,
 		http.MethodPost:    true,
 		http.MethodOptions: true,
-	}, adminAuthHandler(adminServer.HTTPHandler()))))
+	}, adminAuthHandler(adminServer.HTTPHandler())))
 
-	http.Handle("/admin/schema", audit.AuditRequestHttp(adminAuthHandler(http.HandlerFunc(func(
+	baseMux.Handle("/admin/schema", adminAuthHandler(http.HandlerFunc(func(
 		w http.ResponseWriter,
 		r *http.Request) {
 		adminSchemaHandler(w, r, adminServer)
-	}))))
+	})))
 
-	http.Handle("/admin/schema/validate", audit.AuditRequestHttp(http.HandlerFunc(func(w http.ResponseWriter,
+	baseMux.HandleFunc("/admin/schema/validate", func(w http.ResponseWriter,
 		r *http.Request) {
 		schema := readRequest(w, r)
 		w.Header().Set("Content-Type", "application/json")
@@ -507,43 +510,43 @@ func setupServer(closer *z.Closer) {
 		w.WriteHeader(http.StatusBadRequest)
 		errs := strings.Split(strings.TrimSpace(err.Error()), "\n")
 		x.SetStatusWithErrors(w, x.ErrorInvalidRequest, errs)
-	})))
+	})
 
-	http.Handle("/admin/shutdown", audit.AuditRequestHttp(allowedMethodsHandler(allowedMethods{http.
+	baseMux.Handle("/admin/shutdown", allowedMethodsHandler(allowedMethods{http.
 		MethodGet: true},
 		adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			shutDownHandler(w, r, adminServer)
-		})))))
+		}))))
 
-	http.Handle("/admin/draining", audit.AuditRequestHttp(allowedMethodsHandler(allowedMethods{
+	baseMux.Handle("/admin/draining", allowedMethodsHandler(allowedMethods{
 		http.MethodPut:  true,
 		http.MethodPost: true,
 	}, adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		drainingHandler(w, r, adminServer)
-	})))))
+	}))))
 
-	http.Handle("/admin/export", audit.AuditRequestHttp(allowedMethodsHandler(
+	baseMux.Handle("/admin/export", allowedMethodsHandler(
 		allowedMethods{http.MethodGet: true},
 		adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			exportHandler(w, r, adminServer)
-		})))))
+		}))))
 
-	http.Handle("/admin/config/cache_mb", audit.AuditRequestHttp(allowedMethodsHandler(allowedMethods{
+	baseMux.Handle("/admin/config/cache_mb", allowedMethodsHandler(allowedMethods{
 		http.MethodGet: true,
 		http.MethodPut: true,
 	}, adminAuthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		memoryLimitHandler(w, r, adminServer)
-	})))))
+	}))))
 
 	addr := fmt.Sprintf("%s:%d", laddr, httpPort())
 	glog.Infof("Bringing up GraphQL HTTP API at %s/graphql", addr)
 	glog.Infof("Bringing up GraphQL HTTP admin API at %s/admin", addr)
 
 	// Add OpenCensus z-pages.
-	zpages.Handle(http.DefaultServeMux, "/z")
+	zpages.Handle(baseMux, "/z")
 
-	http.Handle("/", audit.AuditRequestHttp(http.HandlerFunc(homeHandler)))
-	http.Handle("/ui/keywords", audit.AuditRequestHttp(http.HandlerFunc(keywordHandler)))
+	baseMux.Handle("/", http.HandlerFunc(homeHandler))
+	baseMux.Handle("/ui/keywords", http.HandlerFunc(keywordHandler))
 
 	// Initialize the servers.
 	admin.ServerCloser.AddRunning(3)
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index dff4bebdfad..585ea315e82 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -283,15 +283,17 @@ func run() {
 	x.Check(err)
 	go x.StartListenHttpAndHttps(httpListener, tlsCfg, st.zero.closer)
 
-	http.HandleFunc("/health", st.pingResponse)
-	http.Handle("/state", audit.AuditRequestHttp(http.HandlerFunc(st.getState)))
-	http.Handle("/removeNode", audit.AuditRequestHttp(http.HandlerFunc(st.removeNode)))
-	http.Handle("/moveTablet", audit.AuditRequestHttp(http.HandlerFunc(st.moveTablet)))
-	http.Handle("/assign", audit.AuditRequestHttp(http.HandlerFunc(st.assign)))
-	http.Handle("/enterpriseLicense",
-		audit.AuditRequestHttp(http.HandlerFunc(st.applyEnterpriseLicense)))
-	http.HandleFunc("/jemalloc", x.JemallocHandler)
-	zpages.Handle(http.DefaultServeMux, "/z")
+	baseMux := http.NewServeMux()
+	http.Handle("/", audit.AuditRequestHttp(baseMux))
+
+	baseMux.HandleFunc("/health", st.pingResponse)
+	baseMux.Handle("/state", http.HandlerFunc(st.getState))
+	baseMux.Handle("/removeNode", http.HandlerFunc(st.removeNode))
+	baseMux.Handle("/moveTablet", http.HandlerFunc(st.moveTablet))
+	baseMux.Handle("/assign", http.HandlerFunc(st.assign))
+	baseMux.Handle("/enterpriseLicense", http.HandlerFunc(st.applyEnterpriseLicense))
+	baseMux.HandleFunc("/jemalloc", x.JemallocHandler)
+	zpages.Handle(baseMux, "/z")
 
 	// This must be here. It does not work if placed before Grpc init.
 	x.Check(st.node.initAndStartNode())
diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
index 597b0243fd5..29187b47c77 100644
--- a/ee/audit/interceptor_ee.go
+++ b/ee/audit/interceptor_ee.go
@@ -65,7 +65,17 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 
 func AuditRequestHttp(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if atomic.LoadUint32(&auditEnabled) == 0 {
+		skip := func(method string) bool {
+			skipApis := map[string]bool{
+				// raft server
+				"/health":   true,
+				"/jemalloc": true,
+				"/state": true,
+			}
+			return skipApis[r.URL.Path]
+		}
+
+		if atomic.LoadUint32(&auditEnabled) == 0 || skip(r.URL.Path) {
 			next.ServeHTTP(w, r)
 			return
 		}

From 512f2914fcec45294aa84700adc4a37bd91369a6 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Wed, 20 Jan 2021 11:29:20 +0530
Subject: [PATCH 22/31] refactoring alpha and audit_ee

---
 dgraph/cmd/alpha/run.go | 18 +++++++++---------
 dgraph/cmd/zero/raft.go |  4 +++-
 dgraph/cmd/zero/run.go  | 10 +++++-----
 ee/audit/audit_ee.go    | 37 +++++++++++++++++++++----------------
 4 files changed, 38 insertions(+), 31 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index d301bcb7b72..2fdfe1d8347 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -436,12 +436,12 @@ func setupServer(closer *z.Closer) {
 	baseMux := http.NewServeMux()
 	http.Handle("/", audit.AuditRequestHttp(baseMux))
 
-	baseMux.Handle("/query", http.HandlerFunc(queryHandler))
-	baseMux.Handle("/query/", http.HandlerFunc(queryHandler))
-	baseMux.Handle("/mutate", http.HandlerFunc(mutationHandler))
-	baseMux.Handle("/mutate/", http.HandlerFunc(mutationHandler))
-	baseMux.Handle("/commit", http.HandlerFunc(commitHandler))
-	baseMux.Handle("/alter", http.HandlerFunc(alterHandler))
+	baseMux.HandleFunc("/query", queryHandler)
+	baseMux.HandleFunc("/query/", queryHandler)
+	baseMux.HandleFunc("/mutate", mutationHandler)
+	baseMux.HandleFunc("/mutate/", mutationHandler)
+	baseMux.HandleFunc("/commit", commitHandler)
+	baseMux.HandleFunc("/alter", alterHandler)
 	baseMux.HandleFunc("/health", healthCheck)
 	baseMux.HandleFunc("/state", stateHandler)
 	baseMux.HandleFunc("/jemalloc", x.JemallocHandler)
@@ -727,6 +727,9 @@ func run() {
 
 	worker.InitServerState()
 
+	// Audit is enterprise feature.
+	x.Check(audit.InitAuditorIfNecessary(opts.Audit, worker.EnterpriseEnabled))
+
 	if Alpha.Conf.GetBool("expose_trace") {
 		// TODO: Remove this once we get rid of event logs.
 		trace.AuthRequest = func(req *http.Request) (any, sensitive bool) {
@@ -807,9 +810,6 @@ func run() {
 	// close alpha. This closer is for closing and waiting that subscription.
 	adminCloser := z.NewCloser(1)
 
-	// Audit is enterprise feature.
-	audit.InitAuditorIfNecessary(opts.Audit, worker.EnterpriseEnabled)
-
 	setupServer(adminCloser)
 	glog.Infoln("GRPC and HTTP stopped.")
 
diff --git a/dgraph/cmd/zero/raft.go b/dgraph/cmd/zero/raft.go
index a4155e55e86..3b38f29c99a 100644
--- a/dgraph/cmd/zero/raft.go
+++ b/dgraph/cmd/zero/raft.go
@@ -406,7 +406,9 @@ func (n *node) applyProposal(e raftpb.Entry) (uint64, error) {
 			if encKey, err := audit.ReadAuditEncKey(opts.audit); err != nil {
 				glog.Errorf("error while reading encryption file %+v", err)
 			} else {
-				audit.InitAuditor(x.GetFlagString(opts.audit, "dir"), encKey)
+				if err := audit.InitAuditor(x.GetFlagString(opts.audit, "dir"), encKey); err != nil {
+					glog.Errorf("error while initializing audit logs %+v", err)
+				}
 			}
 		}
 	}
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index 585ea315e82..b7970ff98f9 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -287,11 +287,11 @@ func run() {
 	http.Handle("/", audit.AuditRequestHttp(baseMux))
 
 	baseMux.HandleFunc("/health", st.pingResponse)
-	baseMux.Handle("/state", http.HandlerFunc(st.getState))
-	baseMux.Handle("/removeNode", http.HandlerFunc(st.removeNode))
-	baseMux.Handle("/moveTablet", http.HandlerFunc(st.moveTablet))
-	baseMux.Handle("/assign", http.HandlerFunc(st.assign))
-	baseMux.Handle("/enterpriseLicense", http.HandlerFunc(st.applyEnterpriseLicense))
+	baseMux.HandleFunc("/state", st.getState)
+	baseMux.HandleFunc("/removeNode", st.removeNode)
+	baseMux.HandleFunc("/moveTablet", st.moveTablet)
+	baseMux.HandleFunc("/assign", st.assign)
+	baseMux.HandleFunc("/enterpriseLicense", st.applyEnterpriseLicense)
 	baseMux.HandleFunc("/jemalloc", x.JemallocHandler)
 	zpages.Handle(baseMux, "/z")
 
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index 011afdc0814..b49335213d3 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -22,6 +22,10 @@ import (
 	"github.com/golang/glog"
 )
 
+const (
+	defaultAuditFilename = "dgraph_audit.log"
+)
+
 var auditEnabled uint32
 
 type AuditEvent struct {
@@ -69,44 +73,43 @@ func ReadAuditEncKey(conf string) ([]byte, error) {
 // InitAuditorIfNecessary accepts conf and enterprise edition check function.
 // This method keep tracks whether cluster is part of enterprise edition or not.
 // It pools eeEnabled function every five minutes to check if the license is still valid or not.
-func InitAuditorIfNecessary(conf string, eeEnabled func() bool) {
+func InitAuditorIfNecessary(conf string, eeEnabled func() bool) error {
 	if conf == "" {
-		return
+		return nil
 	}
 	encKey, err := ReadAuditEncKey(conf)
 	if err != nil {
 		glog.Errorf("error while reading encryption file", err)
-		return
+		return err
 	}
 	if eeEnabled() {
-		InitAuditor(x.GetFlagString(conf, "dir"), encKey)
+		if err := InitAuditor(x.GetFlagString(conf, "dir"), encKey); err != nil {
+			return err
+		}
 	}
 	auditor.tick = time.NewTicker(time.Minute * 5)
 	go trackIfEEValid(x.GetFlagString(conf, "dir"), encKey, eeEnabled)
+	return nil
 }
 
 // InitAuditor initializes the auditor.
 // This method doesnt keep track of whether cluster is part of enterprise edition or not.
 // Client has to keep track of that.
-func InitAuditor(dir string, key []byte) {
-	auditor.log = initlog(dir, key)
+func InitAuditor(dir string, key []byte) error {
+	var err error
+	if auditor.log, err = x.InitLogger(dir, defaultAuditFilename, key); err != nil {
+		return err
+	}
 	atomic.StoreUint32(&auditEnabled, 1)
 	glog.Infoln("audit logs are enabled")
-}
-
-func initlog(dir string, key []byte) *x.Logger {
-	logger, err := x.InitLogger(dir, "dgraph_audit.log", key)
-	if err != nil {
-		glog.Errorf("error while initiating auditor %v", err)
-		return nil
-	}
-	return logger
+	return nil
 }
 
 // trackIfEEValid tracks enterprise license of the cluster.
 // Right now alpha doesn't know about the enterprise/licence.
 // That's why we needed to track if the current node is part of enterprise edition cluster
 func trackIfEEValid(dir string, key []byte, eeEnabledFunc func() bool) {
+	var err error
 	for {
 		select {
 		case <-auditor.tick.C:
@@ -118,7 +121,9 @@ func trackIfEEValid(dir string, key []byte, eeEnabledFunc func() bool) {
 			}
 
 			if atomic.LoadUint32(&auditEnabled) != 1 {
-				auditor.log = initlog(dir, key)
+				if auditor.log, err = x.InitLogger(dir, defaultAuditFilename, key); err != nil {
+					continue
+				}
 				atomic.StoreUint32(&auditEnabled, 1)
 				glog.Infof("audit logs are enabled")
 			}

From 6336d094c8513d4e1838d1b66f68f4dd47c282f2 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Wed, 20 Jan 2021 11:43:48 +0530
Subject: [PATCH 23/31] refactoring some error handling

---
 ee/audit/interceptor_ee.go |  6 +++---
 ee/audit/run_ee.go         |  2 +-
 x/log_writer.go            | 15 +++++----------
 3 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
index 29187b47c77..e8ad205fe13 100644
--- a/ee/audit/interceptor_ee.go
+++ b/ee/audit/interceptor_ee.go
@@ -66,13 +66,13 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 func AuditRequestHttp(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		skip := func(method string) bool {
-			skipApis := map[string]bool{
-				// raft server
+			skipEPs := map[string]bool{
+				// list of endpoints that needs to be skipped
 				"/health":   true,
 				"/jemalloc": true,
 				"/state": true,
 			}
-			return skipApis[r.URL.Path]
+			return skipEPs[r.URL.Path]
 		}
 
 		if atomic.LoadUint32(&auditEnabled) == 0 || skip(r.URL.Path) {
diff --git a/ee/audit/run_ee.go b/ee/audit/run_ee.go
index 08fe2bf97e8..62a9854dd2f 100644
--- a/ee/audit/run_ee.go
+++ b/ee/audit/run_ee.go
@@ -99,7 +99,7 @@ func run() error {
 		stream := cipher.NewCTR(block, iv)
 		stream.XORKeyStream(content, content)
 		x.Check2(outfile.Write(content))
-		// if its the end of data. finish encoding
+		// if its the end of data. finish decrypting
 		if iterator >= stat.Size() {
 			break
 		}
diff --git a/x/log_writer.go b/x/log_writer.go
index d0f9a37f454..23c924f93eb 100644
--- a/x/log_writer.go
+++ b/x/log_writer.go
@@ -60,15 +60,13 @@ func (l *LogWriter) Write(p []byte) (int, error) {
 	if l.file == nil {
 		l.manageLogDir()
 		// open file
-		err = l.open()
-		if err != nil {
+		if err = l.open(); err != nil {
 			return 0, fmt.Errorf("not able to create new file %v", err)
 		}
 	}
 
 	if l.size+int64(len(p)) >= l.MaxSize*1024*1024 {
-		err := l.rotate()
-		if err != nil {
+		if err = l.rotate(); err != nil {
 			return 0, err
 		}
 	}
@@ -100,8 +98,7 @@ func (l *LogWriter) Close() error {
 }
 
 func (l *LogWriter) open() error {
-	err := os.MkdirAll(filepath.Dir(l.FilePath), 0755)
-	if err != nil {
+	if err := os.MkdirAll(filepath.Dir(l.FilePath), 0755); err != nil {
 		return err
 	}
 
@@ -149,8 +146,7 @@ func (l *LogWriter) open() error {
 
 	// If not able to read the baseIv, then this file might be corrupted.
 	// open the new file in that case
-	_, err = f.ReadAt(l.baseIv[:], 0)
-	if err != nil {
+	if _, err = f.ReadAt(l.baseIv[:], 0); err != nil {
 		return openNew()
 	}
 	l.file = f
@@ -169,8 +165,7 @@ func (l *LogWriter) rotate() error {
 		return err
 	}
 
-	_, err = os.Stat(l.FilePath)
-	if err == nil {
+	if _, err = os.Stat(l.FilePath); err != nil {
 		// move the existing file
 		newname := backupName(l.FilePath)
 		if err := os.Rename(l.FilePath, newname); err != nil {

From 5b7c9f475fdfe7b31fc2817a6a6e3352200de383 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Wed, 20 Jan 2021 12:52:17 +0530
Subject: [PATCH 24/31] making audit conf more reliable

---
 dgraph/cmd/alpha/run.go     |  9 ++++---
 dgraph/cmd/zero/raft.go     | 10 +++-----
 dgraph/cmd/zero/run.go      | 18 ++++++-------
 ee/audit/audit.go           | 13 ++++++----
 ee/audit/audit_ee.go        | 51 +++++++++++++++++++++++++++----------
 ee/audit/interceptor.go     |  2 +-
 ee/audit/interceptor_ee.go  |  2 +-
 ee/audit/run_ee.go          | 16 +++++++++---
 ee/utils_ee.go              |  2 +-
 systest/audit/audit_test.go |  6 +++--
 worker/config.go            | 19 +++++---------
 x/config.go                 |  2 +-
 x/logger.go                 |  9 ++++---
 13 files changed, 94 insertions(+), 65 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 2fdfe1d8347..68fbc963f1e 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -608,6 +608,10 @@ func run() {
 	walCache := (cachePercent[3] * (totalCache << 20)) / 100
 
 	ctype, clevel := x.ParseCompression(Alpha.Conf.GetString("badger.compression"))
+
+	x.CheckFlag(Alpha.Conf.GetString("audit"), "dir", "compress", "encrypt-file")
+	conf, err := audit.GetAuditConf(Alpha.Conf.GetString("audit"))
+	x.Check(err)
 	opts := worker.Options{
 		PostingDir:                 Alpha.Conf.GetString("postings"),
 		WALDir:                     Alpha.Conf.GetString("wal"),
@@ -620,10 +624,9 @@ func run() {
 
 		MutationsMode: worker.AllowMutations,
 		AuthToken:     Alpha.Conf.GetString("auth_token"),
-		Audit:         Alpha.Conf.GetString("audit"),
+		Audit:         conf,
 	}
 
-	x.CheckFlag(opts.Audit, "dir", "compress", "encrypt-file")
 	secretFile := Alpha.Conf.GetString("acl_secret_file")
 	if secretFile != "" {
 		hmacSecret, err := ioutil.ReadFile(secretFile)
@@ -684,7 +687,7 @@ func run() {
 		TLSClientConfig:      tlsClientConf,
 		TLSServerConfig:      tlsServerConf,
 		HmacSecret:           opts.HmacSecret,
-		Audit:                opts.Audit,
+		Audit:                opts.Audit != nil,
 	}
 	x.WorkerConfig.Parse(Alpha.Conf)
 	x.CheckFlag(x.WorkerConfig.Raft, "group", "idx", "learner")
diff --git a/dgraph/cmd/zero/raft.go b/dgraph/cmd/zero/raft.go
index 3b38f29c99a..61042574463 100644
--- a/dgraph/cmd/zero/raft.go
+++ b/dgraph/cmd/zero/raft.go
@@ -402,13 +402,9 @@ func (n *node) applyProposal(e raftpb.Entry) (uint64, error) {
 		// Check expiry and set enabled accordingly.
 		expiry := time.Unix(state.License.ExpiryTs, 0).UTC()
 		state.License.Enabled = time.Now().UTC().Before(expiry)
-		if state.License.Enabled && len(opts.audit) > 0 {
-			if encKey, err := audit.ReadAuditEncKey(opts.audit); err != nil {
-				glog.Errorf("error while reading encryption file %+v", err)
-			} else {
-				if err := audit.InitAuditor(x.GetFlagString(opts.audit, "dir"), encKey); err != nil {
-					glog.Errorf("error while initializing audit logs %+v", err)
-				}
+		if state.License.Enabled && opts.audit != nil {
+			if err := audit.InitAuditor(opts.audit); err != nil {
+				glog.Errorf("error while initializing audit logs %+v", err)
 			}
 		}
 	}
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index b7970ff98f9..d2308dc871a 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -57,7 +57,7 @@ type options struct {
 	w                 string
 	rebalanceInterval time.Duration
 	tlsClientConfig   *tls.Config
-	audit             string
+	audit             *audit.AuditConf
 }
 
 var opts options
@@ -199,6 +199,9 @@ func run() {
 	tlsConf, err := x.LoadClientTLSConfigForInternalPort(Zero.Conf)
 	x.Check(err)
 
+	x.CheckFlag(Zero.Conf.GetString("audit"), "dir", "compress", "encrypt-file")
+	conf, err := audit.GetAuditConf(Zero.Conf.GetString("audit"))
+	x.Check(err)
 	opts = options{
 		bindall:           Zero.Conf.GetBool("bindall"),
 		portOffset:        Zero.Conf.GetInt("port_offset"),
@@ -208,12 +211,11 @@ func run() {
 		w:                 Zero.Conf.GetString("wal"),
 		rebalanceInterval: Zero.Conf.GetDuration("rebalance_interval"),
 		tlsClientConfig:   tlsConf,
-		audit:             Zero.Conf.GetString("audit"),
+		audit:             conf,
 	}
 	glog.Infof("Setting Config to: %+v", opts)
 	x.WorkerConfig.Parse(Zero.Conf)
 	x.CheckFlag(opts.raftOpts, "idx", "learner")
-	x.CheckFlag(opts.audit, "dir", "compress", "encrypt-file")
 
 	if !enc.EeBuild && Zero.Conf.GetString("enterprise_license") != "" {
 		log.Fatalf("ERROR: enterprise_license option cannot be applied to OSS builds. ")
@@ -231,16 +233,12 @@ func run() {
 		}
 	}
 
-	if len(opts.audit) > 0 {
+	if opts.audit != nil {
 		wd, err := filepath.Abs(opts.w)
 		x.Check(err)
-		dir := x.GetFlagString(opts.audit, "dir")
-		if len(dir) == 0 {
-			glog.Fatal("audit flag is provided but dir is not specified")
-		}
-		ad, err := filepath.Abs(dir)
+		ad, err := filepath.Abs(opts.audit.Dir)
 		x.Check(err)
-		x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", dir)
+		x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", opts.audit.Dir)
 	}
 
 	if opts.rebalanceInterval <= 0 {
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index bde72a06de2..0d43f1e15f5 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -18,17 +18,20 @@
 
 package audit
 
+type AuditConf struct {
+	Dir string
+}
 
-func ReadAuditEncKey(conf string) ([]byte, error) {
+func GetAuditConf(conf string) (*AuditConf, error) {
 	return nil, nil
 }
 
-func InitAuditorIfNecessary(conf string, eeEnabled func() bool) {
-	return
+func InitAuditorIfNecessary(conf *AuditConf, eeEnabled func() bool) error {
+	return nil
 }
 
-func InitAuditor(dir string, key []byte) {
-	return
+func InitAuditor(conf *AuditConf) error {
+	return nil
 }
 
 func Close() {
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index b49335213d3..86b45b6caf9 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -13,6 +13,7 @@
 package audit
 
 import (
+	"fmt"
 	"io/ioutil"
 	"path/filepath"
 	"sync/atomic"
@@ -28,6 +29,12 @@ const (
 
 var auditEnabled uint32
 
+type AuditConf struct {
+	Compress     bool
+	Dir          string
+	EncryptBytes []byte
+}
+
 type AuditEvent struct {
 	User        string
 	ServerHost  string
@@ -54,7 +61,26 @@ type auditLogger struct {
 	tick *time.Ticker
 }
 
-func ReadAuditEncKey(conf string) ([]byte, error) {
+func GetAuditConf(conf string) (*AuditConf, error) {
+	if conf == "" {
+		return nil, nil
+	}
+	encBytes, err := readAuditEncKey(conf)
+	if err != nil {
+		return nil, err
+	}
+	dir := x.GetFlagString(conf, "dir")
+	if dir == "" {
+		return nil, fmt.Errorf("dir flag is not provided for the audit logs")
+	}
+	return &AuditConf{
+		Compress:     x.GetFlagBool(conf, "compress"),
+		Dir:          dir,
+		EncryptBytes: encBytes,
+	}, nil
+}
+
+func readAuditEncKey(conf string) ([]byte, error) {
 	encFile := x.GetFlagString(conf, "encrypt-file")
 	if encFile == "" {
 		return nil, nil
@@ -73,31 +99,27 @@ func ReadAuditEncKey(conf string) ([]byte, error) {
 // InitAuditorIfNecessary accepts conf and enterprise edition check function.
 // This method keep tracks whether cluster is part of enterprise edition or not.
 // It pools eeEnabled function every five minutes to check if the license is still valid or not.
-func InitAuditorIfNecessary(conf string, eeEnabled func() bool) error {
-	if conf == "" {
+func InitAuditorIfNecessary(conf *AuditConf, eeEnabled func() bool) error {
+	if conf == nil {
 		return nil
 	}
-	encKey, err := ReadAuditEncKey(conf)
-	if err != nil {
-		glog.Errorf("error while reading encryption file", err)
-		return err
-	}
 	if eeEnabled() {
-		if err := InitAuditor(x.GetFlagString(conf, "dir"), encKey); err != nil {
+		if err := InitAuditor(conf); err != nil {
 			return err
 		}
 	}
 	auditor.tick = time.NewTicker(time.Minute * 5)
-	go trackIfEEValid(x.GetFlagString(conf, "dir"), encKey, eeEnabled)
+	go trackIfEEValid(conf, eeEnabled)
 	return nil
 }
 
 // InitAuditor initializes the auditor.
 // This method doesnt keep track of whether cluster is part of enterprise edition or not.
 // Client has to keep track of that.
-func InitAuditor(dir string, key []byte) error {
+func InitAuditor(conf *AuditConf) error {
 	var err error
-	if auditor.log, err = x.InitLogger(dir, defaultAuditFilename, key); err != nil {
+	if auditor.log, err = x.InitLogger(conf.Dir, defaultAuditFilename, conf.EncryptBytes,
+		conf.Compress); err != nil {
 		return err
 	}
 	atomic.StoreUint32(&auditEnabled, 1)
@@ -108,7 +130,7 @@ func InitAuditor(dir string, key []byte) error {
 // trackIfEEValid tracks enterprise license of the cluster.
 // Right now alpha doesn't know about the enterprise/licence.
 // That's why we needed to track if the current node is part of enterprise edition cluster
-func trackIfEEValid(dir string, key []byte, eeEnabledFunc func() bool) {
+func trackIfEEValid(conf *AuditConf, eeEnabledFunc func() bool) {
 	var err error
 	for {
 		select {
@@ -121,7 +143,8 @@ func trackIfEEValid(dir string, key []byte, eeEnabledFunc func() bool) {
 			}
 
 			if atomic.LoadUint32(&auditEnabled) != 1 {
-				if auditor.log, err = x.InitLogger(dir, defaultAuditFilename, key); err != nil {
+				if auditor.log, err = x.InitLogger(conf.Dir, defaultAuditFilename, conf.EncryptBytes,
+					conf.Compress); err != nil {
 					continue
 				}
 				atomic.StoreUint32(&auditEnabled, 1)
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index 49bf0e47940..df6e6a45d9d 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -33,4 +33,4 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		next.ServeHTTP(w, r)
 	})
-}
\ No newline at end of file
+}
diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
index e8ad205fe13..7f5ee66f008 100644
--- a/ee/audit/interceptor_ee.go
+++ b/ee/audit/interceptor_ee.go
@@ -70,7 +70,7 @@ func AuditRequestHttp(next http.Handler) http.Handler {
 				// list of endpoints that needs to be skipped
 				"/health":   true,
 				"/jemalloc": true,
-				"/state": true,
+				"/state":    true,
 			}
 			return skipEPs[r.URL.Path]
 		}
diff --git a/ee/audit/run_ee.go b/ee/audit/run_ee.go
index 62a9854dd2f..8eb0d9d1bec 100644
--- a/ee/audit/run_ee.go
+++ b/ee/audit/run_ee.go
@@ -18,12 +18,14 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"github.com/dgraph-io/dgraph/ee/enc"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+
 	"github.com/dgraph-io/dgraph/x"
 	"github.com/golang/glog"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
-	"os"
 )
 
 var CmdAudit x.SubCommand
@@ -66,13 +68,16 @@ func initSubcommands() []*x.SubCommand {
 	decFlags.String("in", "", "input file that needs to decrypted.")
 	decFlags.String("out", "audit_log_out.log",
 		"output file to which decrypted output will be dumped.")
-	enc.RegisterFlags(decFlags)
+	decFlags.String("encryption_key_file", "", "path to encrypt files.")
 	return []*x.SubCommand{&decryptCmd}
 }
 
 func run() error {
-	key, err := enc.ReadKey(decryptCmd.Conf)
+	path, err := filepath.Abs(decryptCmd.Conf.GetString("encryption_key_file"))
 	x.Check(err)
+	key, err := ioutil.ReadFile(path)
+	x.Check(err)
+
 	if key == nil {
 		return errors.New("no encryption key provided")
 	}
@@ -106,5 +111,8 @@ func run() error {
 		x.Check2(file.ReadAt(iv[12:], iterator))
 		iterator = iterator + 4
 	}
+	glog.Infof("Decryption of Audit file %s is Done. Decrypted file is %s",
+		decryptCmd.Conf.GetString("in"),
+		decryptCmd.Conf.GetString("out"))
 	return nil
 }
diff --git a/ee/utils_ee.go b/ee/utils_ee.go
index d2788dda918..178ff266022 100644
--- a/ee/utils_ee.go
+++ b/ee/utils_ee.go
@@ -37,7 +37,7 @@ func GetEEFeaturesList() []string {
 	} else {
 		ee = append(ee, "backup_restore")
 	}
-	if x.WorkerConfig.Audit != "" {
+	if x.WorkerConfig.Audit {
 		ee = append(ee, "audit")
 	}
 	return ee
diff --git a/systest/audit/audit_test.go b/systest/audit/audit_test.go
index f172460ac54..fe79bcdab1f 100644
--- a/systest/audit/audit_test.go
+++ b/systest/audit/audit_test.go
@@ -20,13 +20,15 @@ import (
 	"bufio"
 	"encoding/json"
 	"fmt"
-	"github.com/dgraph-io/dgraph/testutil"
-	"github.com/stretchr/testify/require"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"testing"
+
+	"github.com/dgraph-io/dgraph/testutil"
+	"github.com/stretchr/testify/require"
 )
+
 func TestZeroAudit(t *testing.T) {
 	zeroCmd := map[string][]string{
 		"/removeNode": []string{`--location`, "--request", "GET",
diff --git a/worker/config.go b/worker/config.go
index 76506305319..d4e22c7aa59 100644
--- a/worker/config.go
+++ b/worker/config.go
@@ -17,12 +17,11 @@
 package worker
 
 import (
-	"github.com/golang/glog"
 	"path/filepath"
 	"time"
 
 	bo "github.com/dgraph-io/badger/v3/options"
-
+	"github.com/dgraph-io/dgraph/ee/audit"
 	"github.com/dgraph-io/dgraph/x"
 )
 
@@ -72,7 +71,7 @@ type Options struct {
 	// CacheMb is the total memory allocated between all the caches.
 	CacheMb int64
 
-	Audit string
+	Audit *audit.AuditConf
 }
 
 // Config holds an instance of the server options..
@@ -100,16 +99,12 @@ func (opt *Options) validate() {
 	x.AssertTruef(pd != wd, "Posting and WAL directory cannot be the same ('%s').", opt.PostingDir)
 	x.AssertTruef(pd != td, "Posting and Tmp directory cannot be the same ('%s').", opt.PostingDir)
 	x.AssertTruef(wd != td, "WAL and Tmp directory cannot be the same ('%s').", opt.WALDir)
-	if opt.Audit !="" {
-		dir := x.GetFlagString(opt.Audit, "dir")
-		if dir == "" {
-			glog.Fatal("audit flag is provided but dir is not specified")
-		}
-		ad, err := filepath.Abs(dir)
+	if opt.Audit != nil {
+		ad, err := filepath.Abs(opt.Audit.Dir)
 		x.Check(err)
-		x.AssertTruef(ad != pd, "Posting and Audit Directory cannot be the same ('%s').", dir)
-		x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", dir)
-		x.AssertTruef(ad != td, "Tmp and Audit directory cannot be the same ('%s').", dir)
+		x.AssertTruef(ad != pd, "Posting and Audit Directory cannot be the same ('%s').", opt.Audit.Dir)
+		x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", opt.Audit.Dir)
+		x.AssertTruef(ad != td, "Tmp and Audit directory cannot be the same ('%s').", opt.Audit.Dir)
 
 	}
 }
diff --git a/x/config.go b/x/config.go
index e4448afc7c7..f6344d081cc 100644
--- a/x/config.go
+++ b/x/config.go
@@ -114,7 +114,7 @@ type WorkerOptions struct {
 	HardSync bool
 
 	// Audit contains the audit flags that enables the audit.
-	Audit string
+	Audit bool
 }
 
 // WorkerConfig stores the global instance of the worker package's options.
diff --git a/x/logger.go b/x/logger.go
index c2927ad5942..0352e9aad5a 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -24,7 +24,7 @@ import (
 	"go.uber.org/zap/zapcore"
 )
 
-func InitLogger(dir string, filename string, key []byte) (*Logger, error) {
+func InitLogger(dir string, filename string, key []byte, compress bool) (*Logger, error) {
 	if err := os.MkdirAll(dir, 0700); err != nil {
 		return nil, err
 	}
@@ -38,10 +38,11 @@ func InitLogger(dir string, filename string, key []byte) (*Logger, error) {
 	}
 	getWriterSyncer := func() zapcore.WriteSyncer {
 		w := &LogWriter{
-			FilePath: path,
-			MaxSize:  100,
-			MaxAge:   30,
+			FilePath:      path,
+			MaxSize:       100,
+			MaxAge:        30,
 			EncryptionKey: key,
+			Compress:      compress,
 		}
 		return zapcore.AddSync(w)
 	}

From 7cd693de073d8e9a589390d80f2ed281daceec7b Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Wed, 20 Jan 2021 13:12:42 +0530
Subject: [PATCH 25/31] fixing goimports

---
 dgraph/cmd/zero/license_ee.go |  3 ++-
 dgraph/cmd/zero/raft.go       |  3 ++-
 ee/audit/interceptor.go       |  3 ++-
 ee/audit/interceptor_ee.go    | 11 ++++++-----
 x/log_writer.go               |  3 ++-
 x/log_writer_test.go          |  3 ++-
 6 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/dgraph/cmd/zero/license_ee.go b/dgraph/cmd/zero/license_ee.go
index e6e7f7c99ce..0da5ed90467 100644
--- a/dgraph/cmd/zero/license_ee.go
+++ b/dgraph/cmd/zero/license_ee.go
@@ -15,12 +15,13 @@ package zero
 import (
 	"bytes"
 	"context"
-	"github.com/dgraph-io/dgraph/ee/audit"
 	"io/ioutil"
 	"math"
 	"net/http"
 	"time"
 
+	"github.com/dgraph-io/dgraph/ee/audit"
+
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/x"
 	"github.com/dgraph-io/ristretto/z"
diff --git a/dgraph/cmd/zero/raft.go b/dgraph/cmd/zero/raft.go
index 61042574463..6444b5acf04 100644
--- a/dgraph/cmd/zero/raft.go
+++ b/dgraph/cmd/zero/raft.go
@@ -20,7 +20,6 @@ import (
 	"context"
 	"encoding/binary"
 	"fmt"
-	"github.com/dgraph-io/dgraph/ee/audit"
 	"log"
 	"math"
 	"sort"
@@ -28,6 +27,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/dgraph-io/dgraph/ee/audit"
+
 	otrace "go.opencensus.io/trace"
 
 	"github.com/dgraph-io/dgraph/conn"
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index df6e6a45d9d..e4d5b2eb081 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -20,8 +20,9 @@ package audit
 
 import (
 	"context"
-	"google.golang.org/grpc"
 	"net/http"
+
+	"google.golang.org/grpc"
 )
 
 func AuditRequestGRPC(ctx context.Context, req interface{},
diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
index 7f5ee66f008..09ed1196659 100644
--- a/ee/audit/interceptor_ee.go
+++ b/ee/audit/interceptor_ee.go
@@ -15,17 +15,18 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"github.com/dgraph-io/dgraph/x"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/metadata"
-	"google.golang.org/grpc/peer"
-	"google.golang.org/grpc/status"
 	"io"
 	"io/ioutil"
 	"net/http"
 	"strings"
 	"sync/atomic"
 
+	"github.com/dgraph-io/dgraph/x"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
+
 	"google.golang.org/grpc"
 )
 
diff --git a/x/log_writer.go b/x/log_writer.go
index 23c924f93eb..f49f4d21df7 100644
--- a/x/log_writer.go
+++ b/x/log_writer.go
@@ -20,7 +20,6 @@ import (
 	"compress/gzip"
 	"encoding/binary"
 	"fmt"
-	"github.com/dgraph-io/badger/v3/y"
 	"io"
 	"io/ioutil"
 	"math/rand"
@@ -29,6 +28,8 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/dgraph-io/badger/v3/y"
 )
 
 const (
diff --git a/x/log_writer_test.go b/x/log_writer_test.go
index 64bbae5b971..d50af393654 100644
--- a/x/log_writer_test.go
+++ b/x/log_writer_test.go
@@ -23,13 +23,14 @@ import (
 	"crypto/aes"
 	"crypto/cipher"
 	"encoding/binary"
-	"github.com/stretchr/testify/require"
 	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strings"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestLogWriter(t *testing.T) {

From 642c296a8f1956cab87042672005144a1af348b4 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Wed, 20 Jan 2021 14:51:48 +0530
Subject: [PATCH 26/31] fixing merge conflicts

---
 dgraph/cmd/alpha/run.go |  1 -
 dgraph/cmd/zero/run.go  |  2 --
 ee/audit/audit_ee.go    | 18 ++++++++++--------
 x/flags.go              |  5 ++---
 4 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index d75c57c04b9..523b7c1acae 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -604,7 +604,6 @@ func run() {
 
 	ctype, clevel := x.ParseCompression(Alpha.Conf.GetString("badger.compression"))
 
-	x.CheckFlag(Alpha.Conf.GetString("audit"), "dir", "compress", "encrypt-file")
 	conf, err := audit.GetAuditConf(Alpha.Conf.GetString("audit"))
 	x.Check(err)
 	opts := worker.Options{
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index c814da1bd56..7ccb68c9d39 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -200,8 +200,6 @@ func run() {
 	x.Check(err)
 
 	raft := x.NewSuperFlag(Zero.Conf.GetString("raft")).MergeAndCheckDefault(raftDefault)
-
-	x.CheckFlag(Zero.Conf.GetString("audit"), "dir", "compress", "encrypt-file")
 	conf, err := audit.GetAuditConf(Zero.Conf.GetString("audit"))
 	x.Check(err)
 	opts = options{
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index 86b45b6caf9..7df717d6327 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -24,6 +24,7 @@ import (
 )
 
 const (
+	defaultAuditConf     = "dir=;compress=false;encrypt-file="
 	defaultAuditFilename = "dgraph_audit.log"
 )
 
@@ -65,23 +66,24 @@ func GetAuditConf(conf string) (*AuditConf, error) {
 	if conf == "" {
 		return nil, nil
 	}
-	encBytes, err := readAuditEncKey(conf)
-	if err != nil {
-		return nil, err
-	}
-	dir := x.GetFlagString(conf, "dir")
+	auditFlag := x.NewSuperFlag(conf).MergeAndCheckDefault(defaultAuditConf)
+	dir := auditFlag.GetString("dir")
 	if dir == "" {
 		return nil, fmt.Errorf("dir flag is not provided for the audit logs")
 	}
+	encBytes, err := readAuditEncKey(auditFlag)
+	if err != nil {
+		return nil, err
+	}
 	return &AuditConf{
-		Compress:     x.GetFlagBool(conf, "compress"),
+		Compress:     auditFlag.GetBool("compress"),
 		Dir:          dir,
 		EncryptBytes: encBytes,
 	}, nil
 }
 
-func readAuditEncKey(conf string) ([]byte, error) {
-	encFile := x.GetFlagString(conf, "encrypt-file")
+func readAuditEncKey(conf *x.SuperFlag) ([]byte, error) {
+	encFile := conf.GetString("encrypt-file")
 	if encFile == "" {
 		return nil, nil
 	}
diff --git a/x/flags.go b/x/flags.go
index 507a2b933e6..aa3cfeb5a93 100644
--- a/x/flags.go
+++ b/x/flags.go
@@ -146,7 +146,6 @@ func (sf *SuperFlag) GetUint32(opt string) uint32 {
 	Checkf(err, "Unable to parse %s as uint32 for key: %s. Options: %s\n", val, opt, sf)
 	return uint32(u)
 }
-func GetFlagString(opt, key string) string {
-	val := GetFlag(opt, key)
-	return val
+func (sf *SuperFlag) GetString(opt string) string {
+	return sf.Get(opt)
 }

From d97d78dc24f89c439e9bd89b99d04edb224d70d3 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Sun, 24 Jan 2021 19:48:21 +0530
Subject: [PATCH 27/31] making log writer performant using buffered writer

---
 dgraph/cmd/bulk/count_index.go |   2 +-
 dgraph/cmd/bulk/reduce.go      |   2 +-
 ee/audit/audit.go              |   2 +-
 ee/audit/interceptor.go        |   2 +-
 ee/audit/interceptor_ee.go     |  46 ++++----
 ee/audit/run.go                |   2 +-
 ee/audit/run_ee.go             |   9 +-
 x/flags.go                     |  17 ++-
 x/flags_test.go                |   2 +-
 x/log_writer.go                | 191 ++++++++++++++++++---------------
 x/log_writer_test.go           |   9 +-
 x/logger.go                    |  26 +++--
 12 files changed, 161 insertions(+), 149 deletions(-)

diff --git a/dgraph/cmd/bulk/count_index.go b/dgraph/cmd/bulk/count_index.go
index 4377d0005f2..12a3e602d15 100644
--- a/dgraph/cmd/bulk/count_index.go
+++ b/dgraph/cmd/bulk/count_index.go
@@ -153,7 +153,7 @@ func (c *countIndexer) writeIndex(buf *z.Buffer) {
 		encoder = codec.Encoder{BlockSize: 256}
 		pl.Reset()
 
-		// Flush out the buffer.
+		// flush out the buffer.
 		if outBuf.LenNoPadding() > 4<<20 {
 			x.Check(c.writer.Write(outBuf))
 			outBuf.Reset()
diff --git a/dgraph/cmd/bulk/reduce.go b/dgraph/cmd/bulk/reduce.go
index 9036ac271cf..8f173e3ea86 100644
--- a/dgraph/cmd/bulk/reduce.go
+++ b/dgraph/cmd/bulk/reduce.go
@@ -292,7 +292,7 @@ func (r *reducer) writeTmpSplits(ci *countIndexer, wg *sync.WaitGroup) {
 		}
 
 		for i := 0; i < len(kvs.Kv); i += maxSplitBatchLen {
-			// Flush the write batch when the max batch length is reached to prevent the
+			// flush the write batch when the max batch length is reached to prevent the
 			// value log from growing over the allowed limit.
 			if splitBatchLen >= maxSplitBatchLen {
 				x.Check(ci.splitWriter.Flush())
diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index 0d43f1e15f5..99304f13287 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -1,7 +1,7 @@
 // +build oss
 
 /*
- * Copyright 2017-2021 Dgraph Labs, Inc. and Contributors
+ * Copyright 2021 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.
diff --git a/ee/audit/interceptor.go b/ee/audit/interceptor.go
index e4d5b2eb081..e663957434d 100644
--- a/ee/audit/interceptor.go
+++ b/ee/audit/interceptor.go
@@ -1,7 +1,7 @@
 // +build oss
 
 /*
- * Copyright 2017-2021 Dgraph Labs, Inc. and Contributors
+ * Copyright 2021 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.
diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
index 09ed1196659..df24f6743a5 100644
--- a/ee/audit/interceptor_ee.go
+++ b/ee/audit/interceptor_ee.go
@@ -34,25 +34,33 @@ const (
 	maxReqLength = 512 << 10 // 512 KB
 )
 
+var skipApis = map[string]bool{
+	// raft server
+	"Heartbeat":   true,
+	"RaftMessage": true,
+	"JoinCluster": true,
+	"IsPeer":      true,
+	// zero server
+	"StreamMembership": true,
+	"UpdateMembership": true,
+	"Oracle":           true,
+	"Timestamps":       true,
+	"ShouldServe":      true,
+	// health server
+	"Check": true,
+	"Watch": true,
+}
+
+var skipEPs = map[string]bool{
+	// list of endpoints that needs to be skipped
+	"/health":   true,
+	"/jemalloc": true,
+	"/state":    true,
+}
+
 func AuditRequestGRPC(ctx context.Context, req interface{},
 	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 	skip := func(method string) bool {
-		skipApis := map[string]bool{
-			// raft server
-			"Heartbeat":   true,
-			"RaftMessage": true,
-			"JoinCluster": true,
-			"IsPeer":      true,
-			// zero server
-			"StreamMembership": true,
-			"UpdateMembership": true,
-			"Oracle":           true,
-			"Timestamps":       true,
-			"ShouldServe":      true,
-			// health server
-			"Check": true,
-			"Watch": true,
-		}
 		return skipApis[info.FullMethod[strings.LastIndex(info.FullMethod, "/")+1:]]
 	}
 
@@ -67,12 +75,6 @@ func AuditRequestGRPC(ctx context.Context, req interface{},
 func AuditRequestHttp(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		skip := func(method string) bool {
-			skipEPs := map[string]bool{
-				// list of endpoints that needs to be skipped
-				"/health":   true,
-				"/jemalloc": true,
-				"/state":    true,
-			}
 			return skipEPs[r.URL.Path]
 		}
 
diff --git a/ee/audit/run.go b/ee/audit/run.go
index 3dc44b7ac39..78545acb16d 100644
--- a/ee/audit/run.go
+++ b/ee/audit/run.go
@@ -1,7 +1,7 @@
 // +build oss
 
 /*
- * Copyright 2017-2021 Dgraph Labs, Inc. and Contributors
+ * Copyright 2021 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.
diff --git a/ee/audit/run_ee.go b/ee/audit/run_ee.go
index 8eb0d9d1bec..7cf478a4bfe 100644
--- a/ee/audit/run_ee.go
+++ b/ee/audit/run_ee.go
@@ -1,7 +1,7 @@
 // +build !oss
 
 /*
- * Copyright 2121 Dgraph Labs, Inc. and Contributors
+ * Copyright 2021 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,7 +20,6 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
-	"path/filepath"
 
 	"github.com/dgraph-io/dgraph/x"
 	"github.com/golang/glog"
@@ -73,14 +72,12 @@ func initSubcommands() []*x.SubCommand {
 }
 
 func run() error {
-	path, err := filepath.Abs(decryptCmd.Conf.GetString("encryption_key_file"))
+	key, err := ioutil.ReadFile(decryptCmd.Conf.GetString("encryption_key_file"))
 	x.Check(err)
-	key, err := ioutil.ReadFile(path)
-	x.Check(err)
-
 	if key == nil {
 		return errors.New("no encryption key provided")
 	}
+
 	file, err := os.Open(decryptCmd.Conf.GetString("in"))
 	x.Check(err)
 	defer file.Close()
diff --git a/x/flags.go b/x/flags.go
index aa3cfeb5a93..718858cc0af 100644
--- a/x/flags.go
+++ b/x/flags.go
@@ -113,14 +113,8 @@ func (sf *SuperFlag) MergeAndCheckDefault(flag string) *SuperFlag {
 	}
 	return sf
 }
-func (sf *SuperFlag) Get(opt string) string {
-	if sf == nil {
-		return ""
-	}
-	return sf.m[opt]
-}
 func (sf *SuperFlag) GetBool(opt string) bool {
-	val := sf.Get(opt)
+	val := sf.GetString(opt)
 	if val == "" {
 		return false
 	}
@@ -129,7 +123,7 @@ func (sf *SuperFlag) GetBool(opt string) bool {
 	return b
 }
 func (sf *SuperFlag) GetUint64(opt string) uint64 {
-	val := sf.Get(opt)
+	val := sf.GetString(opt)
 	if val == "" {
 		return 0
 	}
@@ -138,7 +132,7 @@ func (sf *SuperFlag) GetUint64(opt string) uint64 {
 	return u
 }
 func (sf *SuperFlag) GetUint32(opt string) uint32 {
-	val := sf.Get(opt)
+	val := sf.GetString(opt)
 	if val == "" {
 		return 0
 	}
@@ -147,5 +141,8 @@ func (sf *SuperFlag) GetUint32(opt string) uint32 {
 	return uint32(u)
 }
 func (sf *SuperFlag) GetString(opt string) string {
-	return sf.Get(opt)
+	if sf == nil {
+		return ""
+	}
+	return sf.m[opt]
 }
diff --git a/x/flags_test.go b/x/flags_test.go
index 13e2b8c8a78..d126310a44b 100644
--- a/x/flags_test.go
+++ b/x/flags_test.go
@@ -39,6 +39,6 @@ func TestFlag(t *testing.T) {
 	require.Panics(t, c)
 	require.Equal(t, true, sf.GetBool("bool-key"))
 	require.Equal(t, uint64(5), sf.GetUint64("int-key"))
-	require.Equal(t, "value", sf.Get("string-key"))
+	require.Equal(t, "value", sf.GetString("string-key"))
 	require.Equal(t, uint64(5), sf.GetUint64("other-key"))
 }
diff --git a/x/log_writer.go b/x/log_writer.go
index f49f4d21df7..6ff6e19a3d0 100644
--- a/x/log_writer.go
+++ b/x/log_writer.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2021 Dgraph Labs, Inc. and Contributors
+ * Copyright 2021 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.
@@ -17,6 +17,7 @@
 package x
 
 import (
+	"bufio"
 	"compress/gzip"
 	"encoding/binary"
 	"fmt"
@@ -33,7 +34,9 @@ import (
 )
 
 const (
-	backupTimeFormat = "2021-01-19T15-04-05.000"
+	backupTimeFormat = "2006-01-02T15-04-05.000"
+	bufferSize       = 256 * 1024
+	flushInterval    = 30 * time.Second
 )
 
 var _ io.WriteCloser = (*LogWriter)(nil)
@@ -45,29 +48,40 @@ type LogWriter struct {
 	Compress      bool
 	EncryptionKey []byte
 
-	baseIv          [12]byte
-	mu              sync.Mutex
-	size            int64
-	file            *os.File
-	mch             chan bool
-	startDirManager sync.Once
+	baseIv      [12]byte
+	mu          sync.Mutex
+	size        int64
+	file        *os.File
+	writer      *bufio.Writer
+	flushTicker *time.Ticker
+	// To maintain order of manage old logs calls
+	manageChannel chan bool
+}
+
+func (l *LogWriter) Init() (*LogWriter, error) {
+	l.manageOldLogs()
+	if err := l.open(); err != nil {
+		return nil, fmt.Errorf("not able to create new file %v", err)
+	}
+
+	l.manageChannel = make(chan bool, 1)
+	go func() {
+		for range l.manageChannel {
+			l.manageOldLogs()
+		}
+	}()
+
+	l.flushTicker = time.NewTicker(flushInterval)
+	go l.flushPeriodic()
+	return l, nil
 }
 
 func (l *LogWriter) Write(p []byte) (int, error) {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
-	var err error
-	if l.file == nil {
-		l.manageLogDir()
-		// open file
-		if err = l.open(); err != nil {
-			return 0, fmt.Errorf("not able to create new file %v", err)
-		}
-	}
-
 	if l.size+int64(len(p)) >= l.MaxSize*1024*1024 {
-		if err = l.rotate(); err != nil {
+		if err := l.rotate(); err != nil {
 			return 0, err
 		}
 	}
@@ -78,12 +92,12 @@ func (l *LogWriter) Write(p []byte) (int, error) {
 		if err != nil {
 			return 0, err
 		}
-		n, err := l.file.Write(bytes)
+		n, err := l.writer.Write(bytes)
 		l.size = l.size + int64(n)
 		return n, err
 	}
 
-	n, err := l.file.Write(p)
+	n, err := l.writer.Write(p)
 	l.size = l.size + int64(n)
 	return n, err
 }
@@ -94,10 +108,60 @@ func (l *LogWriter) Close() error {
 	if l.file == nil {
 		return nil
 	}
+	close(l.manageChannel)
+	l.flushTicker.Stop()
+	l.flush()
+	_ = l.file.Close()
+	l.writer = nil
 	l.file = nil
 	return nil
 }
 
+// flushPeriodic periodically flushes the log file buffers.
+func (l *LogWriter) flushPeriodic() {
+	for _ = range l.flushTicker.C {
+		l.mu.Lock()
+		l.flush()
+		l.mu.Unlock()
+	}
+}
+
+// LogWriter should be locked while calling this
+func (l *LogWriter) flush() {
+	_ = l.writer.Flush()
+	_ = l.file.Sync()
+}
+
+func encrypt(key []byte, baseIv [12]byte, src []byte) ([]byte, error) {
+	iv := make([]byte, 16)
+	copy(iv, baseIv[:])
+	binary.BigEndian.PutUint32(iv[12:], uint32(len(src)))
+	allocate, err := y.XORBlockAllocate(src, key, iv)
+	if err != nil {
+		return nil, err
+	}
+	allocate = append(iv[12:], allocate...)
+	return allocate, nil
+}
+
+func (l *LogWriter) rotate() error {
+	l.flush()
+	if err := l.file.Close(); err != nil {
+		return err
+	}
+
+	if _, err := os.Stat(l.FilePath); err == nil {
+		// move the existing file
+		newname := backupName(l.FilePath)
+		if err := os.Rename(l.FilePath, newname); err != nil {
+			return fmt.Errorf("can't rename log file: %s", err)
+		}
+	}
+
+	l.manageChannel <- true
+	return l.open()
+}
+
 func (l *LogWriter) open() error {
 	if err := os.MkdirAll(filepath.Dir(l.FilePath), 0755); err != nil {
 		return err
@@ -116,24 +180,23 @@ func (l *LogWriter) open() error {
 		if err != nil {
 			return err
 		}
+		l.file = f
+		l.writer = bufio.NewWriterSize(f, bufferSize)
+
 		if l.EncryptionKey != nil {
 			rand.Read(l.baseIv[:])
-			if _, err = f.Write(l.baseIv[:]); err != nil {
+			if _, err = l.writer.Write(l.baseIv[:]); err != nil {
 				return err
 			}
 		}
-		l.file = f
 		l.size = size()
 		return nil
 	}
 
 	info, err := os.Stat(l.FilePath)
-	if os.IsNotExist(err) {
+	if err != nil { // if any error try to open new log file itself
 		return openNew()
 	}
-	if err != nil {
-		return err
-	}
 
 	// encryption is enabled and file is corrupted as not able to read the IV
 	if l.EncryptionKey != nil && info.Size() < 12 {
@@ -145,40 +208,21 @@ func (l *LogWriter) open() error {
 		return openNew()
 	}
 
-	// If not able to read the baseIv, then this file might be corrupted.
-	// open the new file in that case
-	if _, err = f.ReadAt(l.baseIv[:], 0); err != nil {
-		return openNew()
+	if l.EncryptionKey != nil {
+		// If not able to read the baseIv, then this file might be corrupted.
+		// open the new file in that case
+		if _, err = f.ReadAt(l.baseIv[:], 0); err != nil {
+			_ = f.Close()
+			return openNew()
+		}
 	}
+
 	l.file = f
+	l.writer = bufio.NewWriterSize(f, bufferSize)
 	l.size = size()
 	return nil
 }
 
-func (l *LogWriter) rotate() error {
-	var err error
-	// file not open
-	if l.file == nil {
-		return l.open()
-	}
-
-	if err = l.file.Close(); err != nil {
-		return err
-	}
-
-	if _, err = os.Stat(l.FilePath); err != nil {
-		// move the existing file
-		newname := backupName(l.FilePath)
-		if err := os.Rename(l.FilePath, newname); err != nil {
-			return fmt.Errorf("can't rename log file: %s", err)
-		}
-	}
-
-	err = l.open()
-	l.manageLogDir()
-	return err
-}
-
 func backupName(name string) string {
 	dir := filepath.Dir(name)
 	prefix, ext := prefixAndExt(name)
@@ -186,18 +230,6 @@ func backupName(name string) string {
 	return filepath.Join(dir, fmt.Sprintf("%s-%s%s", prefix, timestamp, ext))
 }
 
-func encrypt(key []byte, baseIv [12]byte, src []byte) ([]byte, error) {
-	iv := make([]byte, 16)
-	copy(iv, baseIv[:])
-	binary.BigEndian.PutUint32(iv[12:], uint32(len(src)))
-	allocate, err := y.XORBlockAllocate(src, key, iv)
-	if err != nil {
-		return nil, err
-	}
-	allocate = append(iv[12:], allocate...)
-	return allocate, nil
-}
-
 func compress(src string) error {
 	f, err := os.Open(src)
 	if err != nil {
@@ -227,22 +259,6 @@ func compress(src string) error {
 	return nil
 }
 
-func (l *LogWriter) manageLogDir() {
-	l.startDirManager.Do(func() {
-		l.mch = make(chan bool, 1)
-		go func() {
-			for range l.mch {
-				l.manageOldLogs()
-			}
-		}()
-	})
-
-	select {
-	case l.mch <- true:
-	default:
-	}
-}
-
 // this should be called in a serial order
 func (l *LogWriter) manageOldLogs() {
 	toRemove, toKeep, err := processOldLogFiles(l.FilePath, l.MaxSize)
@@ -301,15 +317,14 @@ func processOldLogFiles(fp string, maxAge int64) ([]string, []string, error) {
 	cutoff := time.Now().Add(-1 * diff)
 
 	for _, f := range files {
-		if f.IsDir() ||
-			!strings.HasPrefix(f.Name(), defPrefix) ||
-			!strings.HasSuffix(f.Name(), defExt) ||
-			!strings.HasSuffix(f.Name(), defExt+".gz") {
+		if f.IsDir() || // f is directory
+			!strings.HasPrefix(f.Name(), defPrefix) || // f doesnt start with prefix
+			!(strings.HasSuffix(f.Name(), defExt) || strings.HasSuffix(f.Name(), defExt+".gz")) {
 			continue
 		}
 
-		p, e := prefixAndExt(fp)
-		ts, err := time.Parse(backupTimeFormat, f.Name()[len(p):len(f.Name())-len(e)])
+		_, e := prefixAndExt(fp)
+		ts, err := time.Parse(backupTimeFormat, f.Name()[len(defPrefix):len(f.Name())-len(e)])
 		if err != nil {
 			continue
 		}
diff --git a/x/log_writer_test.go b/x/log_writer_test.go
index d50af393654..5f54434b060 100644
--- a/x/log_writer_test.go
+++ b/x/log_writer_test.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2021 Dgraph Labs, Inc. and Contributors
+ * Copyright 2021 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.
@@ -43,6 +43,7 @@ func TestLogWriter(t *testing.T) {
 		Compress: false,
 	}
 
+	lw, _ = lw.Init()
 	writeToLogWriterAndVerify(t, lw, path)
 }
 
@@ -56,13 +57,14 @@ func TestLogWriterWithCompression(t *testing.T) {
 		Compress: true,
 	}
 
+	lw, _ = lw.Init()
 	writeToLogWriterAndVerify(t, lw, path)
 }
 
 // if this test failed and you changed anything, please check the dgraph audit decrypt command.
 // The dgraph audit decrypt command uses the same decryption method
 func TestLogWriterWithEncryption(t *testing.T) {
-	path, _ := filepath.Abs("./log_test/audit.log")
+	path, _ := filepath.Abs("./log_test/audit.log.enc")
 	defer os.RemoveAll(filepath.Dir(path))
 	lw := &LogWriter{
 		FilePath:      path,
@@ -72,10 +74,11 @@ func TestLogWriterWithEncryption(t *testing.T) {
 		EncryptionKey: []byte("1234567890123456"),
 	}
 
+	lw, _ = lw.Init()
 	msg := []byte("abcd")
 	msg = bytes.Repeat(msg, 256)
 	msg[1023] = '\n'
-	for i := 0; i < 100; i++ {
+	for i := 0; i < 10000; i++ {
 		n, err := lw.Write(msg)
 		require.Nil(t, err)
 		require.Equal(t, n, len(msg)+4, "write length is not equal")
diff --git a/x/logger.go b/x/logger.go
index 0352e9aad5a..ef51b6db87f 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2021 Dgraph Labs, Inc. and Contributors
+ * Copyright 2021 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.
@@ -36,21 +36,19 @@ func InitLogger(dir string, filename string, key []byte, compress bool) (*Logger
 	if err != nil {
 		return nil, err
 	}
-	getWriterSyncer := func() zapcore.WriteSyncer {
-		w := &LogWriter{
-			FilePath:      path,
-			MaxSize:       100,
-			MaxAge:        30,
-			EncryptionKey: key,
-			Compress:      compress,
-		}
-		return zapcore.AddSync(w)
+	w := &LogWriter{
+		FilePath:      path,
+		MaxSize:       100,
+		MaxAge:        30,
+		EncryptionKey: key,
+		Compress:      compress,
+	}
+	if w, err = w.Init(); err != nil {
+		return nil, err
 	}
-	core := zapcore.NewCore(zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
-		getWriterSyncer(), zap.DebugLevel)
-
 	return &Logger{
-		logger: zap.New(core),
+		logger: zap.New(zapcore.NewCore(zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
+			zapcore.AddSync(w), zap.DebugLevel)),
 	}, nil
 }
 

From 05ea358d5e970a91365b98aa47359aba4363fd55 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Mon, 25 Jan 2021 01:31:44 +0530
Subject: [PATCH 28/31] fixing zero test case and comments

---
 dgraph/cmd/alpha/run.go     |  2 +-
 dgraph/cmd/zero/run.go      |  2 +-
 ee/audit/audit_ee.go        |  4 ++--
 ee/audit/interceptor_ee.go  |  3 ++-
 ee/audit/run_ee.go          |  7 ++++---
 systest/audit/audit_test.go | 32 ++++++++++++++++++++------------
 worker/config.go            | 19 ++++++++++++-------
 x/log_writer.go             |  1 +
 x/logger.go                 |  3 +++
 9 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 523b7c1acae..682e7140d41 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -198,7 +198,7 @@ they form a Raft group and provide synchronous replication.
 	dir=/path/to/audits to define the path where to store the audit logs.
 	compress=true/false to enabled the compression of old audit logs (default behaviour is false).
 	encrypt_file=enc/key/file enables the audit log encryption with the key path provided with the
-	flag. 
+	flag.
 	Sample flag could look like --audit dir=aa;encrypt_file=/filepath;compress=true`)
 
 	// TLS configurations
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index 7ccb68c9d39..cc923830905 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -105,7 +105,7 @@ instances to achieve high-availability.
 	dir=/path/to/audits to define the path where to store the audit logs.
 	compress=true/false to enabled the compression of old audit logs (default behaviour is false).
 	encrypt_file=enc/key/file enables the audit log encryption with the key path provided with the
-	flag. 
+	flag.
 	Sample flag could look like --audit dir=aa;encrypt_file=/filepath;compress=true`)
 
 	// TLS configurations
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index 7df717d6327..6bf07bd3761 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -145,8 +145,8 @@ func trackIfEEValid(conf *AuditConf, eeEnabledFunc func() bool) {
 			}
 
 			if atomic.LoadUint32(&auditEnabled) != 1 {
-				if auditor.log, err = x.InitLogger(conf.Dir, defaultAuditFilename, conf.EncryptBytes,
-					conf.Compress); err != nil {
+				if auditor.log, err = x.InitLogger(conf.Dir, defaultAuditFilename,
+					conf.EncryptBytes, conf.Compress); err != nil {
 					continue
 				}
 				atomic.StoreUint32(&auditEnabled, 1)
diff --git a/ee/audit/interceptor_ee.go b/ee/audit/interceptor_ee.go
index df24f6743a5..25eb6e5df1e 100644
--- a/ee/audit/interceptor_ee.go
+++ b/ee/audit/interceptor_ee.go
@@ -31,7 +31,7 @@ import (
 )
 
 const (
-	maxReqLength = 512 << 10 // 512 KB
+	maxReqLength = 4 << 10 // 4 KB
 )
 
 var skipApis = map[string]bool{
@@ -46,6 +46,7 @@ var skipApis = map[string]bool{
 	"Oracle":           true,
 	"Timestamps":       true,
 	"ShouldServe":      true,
+	"Connect":          true,
 	// health server
 	"Check": true,
 	"Watch": true,
diff --git a/ee/audit/run_ee.go b/ee/audit/run_ee.go
index 7cf478a4bfe..a4d9ee5734c 100644
--- a/ee/audit/run_ee.go
+++ b/ee/audit/run_ee.go
@@ -43,7 +43,8 @@ func init() {
 			glog.Fatalf("Unable to bind flags for command %v: %v", sc, err)
 		}
 		if err := sc.Conf.BindPFlags(CmdAudit.Cmd.PersistentFlags()); err != nil {
-			glog.Fatalf("Unable to bind persistent flags from audit for command %v: %v", sc, err)
+			glog.Fatalf(
+				"Unable to bind persistent flags from audit for command %v: %v", sc, err)
 		}
 		sc.Conf.SetEnvPrefix(sc.EnvPrefix)
 	}
@@ -82,8 +83,8 @@ func run() error {
 	x.Check(err)
 	defer file.Close()
 
-	outfile, err := os.OpenFile(decryptCmd.Conf.GetString("out"), os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
-		os.ModePerm)
+	outfile, err := os.OpenFile(decryptCmd.Conf.GetString("out"),
+		os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
 	x.Check(err)
 	defer outfile.Close()
 
diff --git a/systest/audit/audit_test.go b/systest/audit/audit_test.go
index fe79bcdab1f..eec0ea57c3a 100644
--- a/systest/audit/audit_test.go
+++ b/systest/audit/audit_test.go
@@ -30,6 +30,7 @@ import (
 )
 
 func TestZeroAudit(t *testing.T) {
+	defer os.RemoveAll("audit_dir/za/dgraph_audit.log")
 	zeroCmd := map[string][]string{
 		"/removeNode": []string{`--location`, "--request", "GET",
 			fmt.Sprintf("%s/removeNode?id=3&group=1", testutil.SockAddrZeroHttp)},
@@ -39,18 +40,22 @@ func TestZeroAudit(t *testing.T) {
 			fmt.Sprintf("%s/moveTablet?tablet=name&group=2", testutil.SockAddrZeroHttp)}}
 
 	msgs := make([]string, 0)
-	for req, c := range zeroCmd {
-		msgs = append(msgs, req)
-		cmd := exec.Command("curl", c...)
-		if out, err := cmd.CombinedOutput(); err != nil {
-			fmt.Println(string(out))
-			t.Fatal(err)
+	// logger is buffered. make calls in bunch so that dont want to wait for flush
+	for i := 0; i < 500; i++ {
+		for req, c := range zeroCmd {
+			msgs = append(msgs, req)
+			cmd := exec.Command("curl", c...)
+			if out, err := cmd.CombinedOutput(); err != nil {
+				fmt.Println(string(out))
+				t.Fatal(err)
+			}
 		}
 	}
 
 	verifyLogs(t, "./audit_dir/za/dgraph_audit.log", msgs)
 }
 func TestAlphaAudit(t *testing.T) {
+	defer os.Remove("audit_dir/aa/dgraph_audit.log")
 	testCommand := map[string][]string{
 		"/admin": []string{"--location", "--request", "POST",
 			fmt.Sprintf("%s/admin", testutil.SockAddrHttp),
@@ -90,12 +95,15 @@ input: {destination: \"/Users/sankalanparajuli/work/backup\"}) {\n    response {
 	}
 
 	msgs := make([]string, 0)
-	for req, c := range testCommand {
-		msgs = append(msgs, req)
-		cmd := exec.Command("curl", c...)
-		if out, err := cmd.CombinedOutput(); err != nil {
-			fmt.Println(string(out))
-			t.Fatal(err)
+	// logger is buffered. make calls in bunch so that dont want to wait for flush
+	for i := 0; i < 200; i++ {
+		for req, c := range testCommand {
+			msgs = append(msgs, req)
+			cmd := exec.Command("curl", c...)
+			if out, err := cmd.CombinedOutput(); err != nil {
+				fmt.Println(string(out))
+				t.Fatal(err)
+			}
 		}
 	}
 	verifyLogs(t, "./audit_dir/aa/dgraph_audit.log", msgs)
diff --git a/worker/config.go b/worker/config.go
index d4e22c7aa59..67eef58614c 100644
--- a/worker/config.go
+++ b/worker/config.go
@@ -96,15 +96,20 @@ func (opt *Options) validate() {
 	x.Check(err)
 	td, err := filepath.Abs(x.WorkerConfig.TmpDir)
 	x.Check(err)
-	x.AssertTruef(pd != wd, "Posting and WAL directory cannot be the same ('%s').", opt.PostingDir)
-	x.AssertTruef(pd != td, "Posting and Tmp directory cannot be the same ('%s').", opt.PostingDir)
-	x.AssertTruef(wd != td, "WAL and Tmp directory cannot be the same ('%s').", opt.WALDir)
+	x.AssertTruef(pd != wd,
+		"Posting and WAL directory cannot be the same ('%s').", opt.PostingDir)
+	x.AssertTruef(pd != td,
+		"Posting and Tmp directory cannot be the same ('%s').", opt.PostingDir)
+	x.AssertTruef(wd != td,
+		"WAL and Tmp directory cannot be the same ('%s').", opt.WALDir)
 	if opt.Audit != nil {
 		ad, err := filepath.Abs(opt.Audit.Dir)
 		x.Check(err)
-		x.AssertTruef(ad != pd, "Posting and Audit Directory cannot be the same ('%s').", opt.Audit.Dir)
-		x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", opt.Audit.Dir)
-		x.AssertTruef(ad != td, "Tmp and Audit directory cannot be the same ('%s').", opt.Audit.Dir)
-
+		x.AssertTruef(ad != pd,
+			"Posting and Audit Directory cannot be the same ('%s').", opt.Audit.Dir)
+		x.AssertTruef(ad != wd,
+			"WAL and Audit directory cannot be the same ('%s').", opt.Audit.Dir)
+		x.AssertTruef(ad != td,
+			"Tmp and Audit directory cannot be the same ('%s').", opt.Audit.Dir)
 	}
 }
diff --git a/x/log_writer.go b/x/log_writer.go
index 6ff6e19a3d0..4c705110a9d 100644
--- a/x/log_writer.go
+++ b/x/log_writer.go
@@ -39,6 +39,7 @@ const (
 	flushInterval    = 30 * time.Second
 )
 
+// This is done to ensure LogWriter always implement io.WriterCloser
 var _ io.WriteCloser = (*LogWriter)(nil)
 
 type LogWriter struct {
diff --git a/x/logger.go b/x/logger.go
index ef51b6db87f..9cab3693656 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -49,11 +49,13 @@ func InitLogger(dir string, filename string, key []byte, compress bool) (*Logger
 	return &Logger{
 		logger: zap.New(zapcore.NewCore(zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
 			zapcore.AddSync(w), zap.DebugLevel)),
+		writer: w,
 	}, nil
 }
 
 type Logger struct {
 	logger *zap.Logger
+	writer *LogWriter
 }
 
 // AuditI logs audit message as info. args are key value pairs with key as string value
@@ -84,4 +86,5 @@ func (l *Logger) Sync() {
 		return
 	}
 	_ = l.logger.Sync()
+	_ = l.writer.Close()
 }

From d01d743403108f4ffe31cde2a123575cfb882dc7 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Tue, 26 Jan 2021 13:52:10 +0530
Subject: [PATCH 29/31] gracefully closing all the go routines

---
 dgraph/cmd/alpha/run.go       |  5 ++---
 dgraph/cmd/zero/license_ee.go |  1 -
 dgraph/cmd/zero/run.go        | 11 ++++++-----
 ee/audit/audit_ee.go          | 35 ++++++++++++++++++++++-------------
 x/log_writer.go               | 34 +++++++++++++++++++++++++---------
 x/logger.go                   |  2 +-
 6 files changed, 56 insertions(+), 32 deletions(-)

diff --git a/dgraph/cmd/alpha/run.go b/dgraph/cmd/alpha/run.go
index 682e7140d41..4573e9eada8 100644
--- a/dgraph/cmd/alpha/run.go
+++ b/dgraph/cmd/alpha/run.go
@@ -604,8 +604,7 @@ func run() {
 
 	ctype, clevel := x.ParseCompression(Alpha.Conf.GetString("badger.compression"))
 
-	conf, err := audit.GetAuditConf(Alpha.Conf.GetString("audit"))
-	x.Check(err)
+	conf := audit.GetAuditConf(Alpha.Conf.GetString("audit"))
 	opts := worker.Options{
 		PostingDir:                 Alpha.Conf.GetString("postings"),
 		WALDir:                     Alpha.Conf.GetString("wal"),
@@ -820,7 +819,7 @@ func run() {
 	glog.Infoln("adminCloser closed.")
 
 	audit.Close()
-	glog.Infoln("audit logs if enabled are closed.")
+
 	worker.State.Dispose()
 	x.RemoveCidFile()
 	glog.Info("worker.State disposed.")
diff --git a/dgraph/cmd/zero/license_ee.go b/dgraph/cmd/zero/license_ee.go
index 0da5ed90467..c9c714b14db 100644
--- a/dgraph/cmd/zero/license_ee.go
+++ b/dgraph/cmd/zero/license_ee.go
@@ -94,7 +94,6 @@ func (n *node) updateEnterpriseState(closer *z.Closer) {
 			if !active {
 				n.server.expireLicense()
 				audit.Close()
-				glog.Infoln("audit logs if enabled are closed.")
 				glog.Warningf("Your enterprise license has expired and enterprise features are " +
 					"disabled. To continue using enterprise features, apply a valid license. To receive " +
 					"a new license, contact us at https://dgraph.io/contact.")
diff --git a/dgraph/cmd/zero/run.go b/dgraph/cmd/zero/run.go
index cc923830905..c6b09383a50 100644
--- a/dgraph/cmd/zero/run.go
+++ b/dgraph/cmd/zero/run.go
@@ -200,8 +200,7 @@ func run() {
 	x.Check(err)
 
 	raft := x.NewSuperFlag(Zero.Conf.GetString("raft")).MergeAndCheckDefault(raftDefault)
-	conf, err := audit.GetAuditConf(Zero.Conf.GetString("audit"))
-	x.Check(err)
+	conf := audit.GetAuditConf(Zero.Conf.GetString("audit"))
 	opts = options{
 		bindall:           Zero.Conf.GetBool("bindall"),
 		portOffset:        Zero.Conf.GetInt("port_offset"),
@@ -237,7 +236,8 @@ func run() {
 		x.Check(err)
 		ad, err := filepath.Abs(opts.audit.Dir)
 		x.Check(err)
-		x.AssertTruef(ad != wd, "WAL and Audit directory cannot be the same ('%s').", opts.audit.Dir)
+		x.AssertTruef(ad != wd,
+			"WAL and Audit directory cannot be the same ('%s').", opts.audit.Dir)
 	}
 
 	if opts.rebalanceInterval <= 0 {
@@ -348,8 +348,9 @@ func run() {
 
 	err = store.Close()
 	glog.Infof("Raft WAL closed with err: %v\n", err)
+
+	audit.Close()
+
 	st.zero.orc.close()
 	glog.Infoln("All done. Goodbye!")
-	audit.Close()
-	glog.Infoln("audit logs if enabled are closed.")
 }
diff --git a/ee/audit/audit_ee.go b/ee/audit/audit_ee.go
index 6bf07bd3761..890f4c9b2d5 100644
--- a/ee/audit/audit_ee.go
+++ b/ee/audit/audit_ee.go
@@ -13,18 +13,19 @@
 package audit
 
 import (
-	"fmt"
 	"io/ioutil"
 	"path/filepath"
 	"sync/atomic"
 	"time"
 
+	"github.com/dgraph-io/ristretto/z"
+
 	"github.com/dgraph-io/dgraph/x"
 	"github.com/golang/glog"
 )
 
 const (
-	defaultAuditConf     = "dir=;compress=false;encrypt-file="
+	defaultAuditConf     = "dir=; compress=false; encrypt-file="
 	defaultAuditFilename = "dgraph_audit.log"
 )
 
@@ -58,28 +59,25 @@ const (
 var auditor *auditLogger = &auditLogger{}
 
 type auditLogger struct {
-	log  *x.Logger
-	tick *time.Ticker
+	log    *x.Logger
+	tick   *time.Ticker
+	closer *z.Closer
 }
 
-func GetAuditConf(conf string) (*AuditConf, error) {
+func GetAuditConf(conf string) *AuditConf {
 	if conf == "" {
-		return nil, nil
+		return nil
 	}
 	auditFlag := x.NewSuperFlag(conf).MergeAndCheckDefault(defaultAuditConf)
 	dir := auditFlag.GetString("dir")
-	if dir == "" {
-		return nil, fmt.Errorf("dir flag is not provided for the audit logs")
-	}
+	x.AssertTruef(dir != "", "dir flag is not provided for the audit logs")
 	encBytes, err := readAuditEncKey(auditFlag)
-	if err != nil {
-		return nil, err
-	}
+	x.Check(err)
 	return &AuditConf{
 		Compress:     auditFlag.GetBool("compress"),
 		Dir:          dir,
 		EncryptBytes: encBytes,
-	}, nil
+	}
 }
 
 func readAuditEncKey(conf *x.SuperFlag) ([]byte, error) {
@@ -111,6 +109,7 @@ func InitAuditorIfNecessary(conf *AuditConf, eeEnabled func() bool) error {
 		}
 	}
 	auditor.tick = time.NewTicker(time.Minute * 5)
+	auditor.closer = z.NewCloser(1)
 	go trackIfEEValid(conf, eeEnabled)
 	return nil
 }
@@ -133,6 +132,7 @@ func InitAuditor(conf *AuditConf) error {
 // Right now alpha doesn't know about the enterprise/licence.
 // That's why we needed to track if the current node is part of enterprise edition cluster
 func trackIfEEValid(conf *AuditConf, eeEnabledFunc func() bool) {
+	defer auditor.closer.Done()
 	var err error
 	for {
 		select {
@@ -152,6 +152,8 @@ func trackIfEEValid(conf *AuditConf, eeEnabledFunc func() bool) {
 				atomic.StoreUint32(&auditEnabled, 1)
 				glog.Infof("audit logs are enabled")
 			}
+		case <-auditor.closer.HasBeenClosed():
+			return
 		}
 	}
 }
@@ -160,11 +162,18 @@ func trackIfEEValid(conf *AuditConf, eeEnabledFunc func() bool) {
 // It also sets the log to nil, because its being called by zero when license expires.
 // If license added, InitLogger will take care of the file.
 func Close() {
+	if atomic.LoadUint32(&auditEnabled) == 0 {
+		return
+	}
 	if auditor.tick != nil {
 		auditor.tick.Stop()
 	}
+	if auditor.closer != nil {
+		auditor.closer.SignalAndWait()
+	}
 	auditor.log.Sync()
 	auditor.log = nil
+	glog.Infoln("audit logs are closed.")
 }
 
 func (a *auditLogger) Audit(event *AuditEvent) {
diff --git a/x/log_writer.go b/x/log_writer.go
index 4c705110a9d..58b8591808d 100644
--- a/x/log_writer.go
+++ b/x/log_writer.go
@@ -30,13 +30,15 @@ import (
 	"sync"
 	"time"
 
+	"github.com/dgraph-io/ristretto/z"
+
 	"github.com/dgraph-io/badger/v3/y"
 )
 
 const (
 	backupTimeFormat = "2006-01-02T15-04-05.000"
 	bufferSize       = 256 * 1024
-	flushInterval    = 30 * time.Second
+	flushInterval    = 10 * time.Second
 )
 
 // This is done to ensure LogWriter always implement io.WriterCloser
@@ -55,7 +57,8 @@ type LogWriter struct {
 	file        *os.File
 	writer      *bufio.Writer
 	flushTicker *time.Ticker
-	// To maintain order of manage old logs calls
+	closer      *z.Closer
+	// To manage order of cleaning old logs files
 	manageChannel chan bool
 }
 
@@ -64,11 +67,17 @@ func (l *LogWriter) Init() (*LogWriter, error) {
 	if err := l.open(); err != nil {
 		return nil, fmt.Errorf("not able to create new file %v", err)
 	}
-
+	l.closer = z.NewCloser(2)
 	l.manageChannel = make(chan bool, 1)
 	go func() {
-		for range l.manageChannel {
-			l.manageOldLogs()
+		defer l.closer.Done()
+		for {
+			select {
+			case <-l.manageChannel:
+				l.manageOldLogs()
+			case <-l.closer.HasBeenClosed():
+				return
+			}
 		}
 	}()
 
@@ -112,6 +121,7 @@ func (l *LogWriter) Close() error {
 	close(l.manageChannel)
 	l.flushTicker.Stop()
 	l.flush()
+	l.closer.SignalAndWait()
 	_ = l.file.Close()
 	l.writer = nil
 	l.file = nil
@@ -120,10 +130,16 @@ func (l *LogWriter) Close() error {
 
 // flushPeriodic periodically flushes the log file buffers.
 func (l *LogWriter) flushPeriodic() {
-	for _ = range l.flushTicker.C {
-		l.mu.Lock()
-		l.flush()
-		l.mu.Unlock()
+	defer l.closer.Done()
+	for {
+		select {
+		case <-l.flushTicker.C:
+			l.mu.Lock()
+			l.flush()
+			l.mu.Unlock()
+		case <-l.closer.HasBeenClosed():
+			return
+		}
 	}
 }
 
diff --git a/x/logger.go b/x/logger.go
index 9cab3693656..08610805322 100644
--- a/x/logger.go
+++ b/x/logger.go
@@ -39,7 +39,7 @@ func InitLogger(dir string, filename string, key []byte, compress bool) (*Logger
 	w := &LogWriter{
 		FilePath:      path,
 		MaxSize:       100,
-		MaxAge:        30,
+		MaxAge:        10,
 		EncryptionKey: key,
 		Compress:      compress,
 	}

From 2498f6d42178686b867c4e6c0e93101117d913e3 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Tue, 26 Jan 2021 13:57:34 +0530
Subject: [PATCH 30/31] fixing oss build

---
 ee/audit/audit.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ee/audit/audit.go b/ee/audit/audit.go
index 99304f13287..772259331a3 100644
--- a/ee/audit/audit.go
+++ b/ee/audit/audit.go
@@ -22,8 +22,8 @@ type AuditConf struct {
 	Dir string
 }
 
-func GetAuditConf(conf string) (*AuditConf, error) {
-	return nil, nil
+func GetAuditConf(conf string) *AuditConf {
+	return nil
 }
 
 func InitAuditorIfNecessary(conf *AuditConf, eeEnabled func() bool) error {

From e35918cb02ad4977b6e3f1ba3ddd9d88043ef464 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Tue, 26 Jan 2021 16:45:50 +0530
Subject: [PATCH 31/31] fixing failed test case

---
 x/log_writer.go | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/x/log_writer.go b/x/log_writer.go
index 58b8591808d..690074f59cc 100644
--- a/x/log_writer.go
+++ b/x/log_writer.go
@@ -113,15 +113,17 @@ func (l *LogWriter) Write(p []byte) (int, error) {
 }
 
 func (l *LogWriter) Close() error {
+	// close all go routines first before acquiring the lock to avoid contention
+	l.closer.SignalAndWait()
+
 	l.mu.Lock()
 	defer l.mu.Unlock()
 	if l.file == nil {
 		return nil
 	}
-	close(l.manageChannel)
-	l.flushTicker.Stop()
 	l.flush()
-	l.closer.SignalAndWait()
+	l.flushTicker.Stop()
+	close(l.manageChannel)
 	_ = l.file.Close()
 	l.writer = nil
 	l.file = nil