Fix for out of memory error with request body

v3/integrations/nrgrpc/go.mod

@@ -13,27 +13,5 @@ require (
 	google.golang.org/protobuf v1.30.0
 )
 
-require (
-	github.com/andybalholm/brotli v1.0.5 // indirect
-	github.com/dlclark/regexp2 v1.9.0 // indirect
-	github.com/gorilla/websocket v1.5.0 // indirect
-	github.com/juju/fslock v0.0.0-20160525022230-4d5c94c67b4b // indirect
-	github.com/k2io/hookingo v1.0.3 // indirect
-	github.com/klauspost/compress v1.16.3 // indirect
-	github.com/kr/pretty v0.3.1 // indirect
-	github.com/mackerelio/go-osstat v0.2.4 // indirect
-	github.com/newrelic/csec-go-agent v0.3.0 // indirect
-	github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect
-	github.com/sirupsen/logrus v1.9.0 // indirect
-	github.com/struCoder/pidusage v0.2.1 // indirect
-	github.com/valyala/bytebufferpool v1.0.0 // indirect
-	github.com/valyala/fasthttp v1.49.0 // indirect
-	golang.org/x/arch v0.3.0 // indirect
-	golang.org/x/net v0.9.0 // indirect
-	golang.org/x/sys v0.7.0 // indirect
-	golang.org/x/text v0.9.0 // indirect
-	google.golang.org/genproto v0.0.0-20230410155749-daa745c078e1 // indirect
-	gopkg.in/yaml.v2 v2.4.0 // indirect
-)
-
 replace github.com/newrelic/go-agent/v3 => ../..
+replace github.com/newrelic/go-agent/v3/integrations/nrsecurityagent => ../../integrations/nrsecurityagent

v3/integrations/nrmicro/nrmicro.go

@@ -5,6 +5,7 @@ package nrmicro
 
 import (
 	"context"
+	"io"
 	"net/http"
 	"net/url"
 	"strings"
@@ -15,9 +16,11 @@ import (
 	"github.com/micro/go-micro/registry"
 	"github.com/micro/go-micro/server"
 
+	protoV1 "github.com/golang/protobuf/proto"
 	"github.com/newrelic/go-agent/v3/internal"
 	"github.com/newrelic/go-agent/v3/internal/integrationsupport"
 	"github.com/newrelic/go-agent/v3/newrelic"
+	protoV2 "google.golang.org/protobuf/proto"
 )
 
 type nrWrapper struct {
@@ -162,7 +165,19 @@ func HandlerWrapper(app *newrelic.Application) server.HandlerWrapper {
 		return func(ctx context.Context, req server.Request, rsp interface{}) error {
 			txn := startWebTransaction(ctx, app, req)
 			defer txn.End()
-			err := fn(newrelic.NewContext(ctx, txn), req, rsp)
+			if req.Body() != nil && newrelic.IsSecurityAgentPresent() {
+				messageType, version := getMessageType(req.Body())
+				newrelic.GetSecurityAgentInterface().SendEvent("GRPC", req.Body(), messageType, version)
+			}
+
+			nrrsp := rsp
+			if req.Stream() && newrelic.IsSecurityAgentPresent() {
+				if stream, ok := rsp.(server.Stream); ok {
+					nrrsp = wrappedServerStream{stream}
+				}
+			}
+
+			err := fn(newrelic.NewContext(ctx, txn), req, nrrsp)
 			var code int
 			if err != nil {
 				if t, ok := err.(*errors.Error); ok {
@@ -227,9 +242,6 @@ func SubscriberWrapper(app *newrelic.Application) server.SubscriberWrapper {
 
 func startWebTransaction(ctx context.Context, app *newrelic.Application, req server.Request) *newrelic.Transaction {
 	var hdrs http.Header
-	var unencodedBody []byte
-	var err error
-
 	if md, ok := metadata.FromContext(ctx); ok {
 		hdrs = make(http.Header, len(md))
 		for k, v := range md {
@@ -242,20 +254,58 @@ func startWebTransaction(ctx context.Context, app *newrelic.Application, req ser
 		Host:   req.Service(),
 		Path:   req.Endpoint(),
 	}
-
-	if unencodedBody, err = req.Read(); err != nil {
-		unencodedBody = nil
-	}
-
 	webReq := newrelic.WebRequest{
 		Header:    hdrs,
 		URL:       u,
 		Method:    req.Method(),
 		Transport: newrelic.TransportHTTP,
-		Body:      unencodedBody,
-		Type:      "HTTP",
+		Type:      "micro",
 	}
 	txn.SetWebRequest(webReq)
 
 	return txn
 }
+
+type wrappedServerStream struct {
+	stream server.Stream
+}
+
+func (s wrappedServerStream) Context() context.Context {
+	return s.stream.Context()
+}
+func (s wrappedServerStream) Request() server.Request {
+	return s.stream.Request()
+}
+func (s wrappedServerStream) Send(msg any) error {
+	return s.stream.Send(msg)
+}
+func (s wrappedServerStream) Recv(msg any) error {
+	err := s.stream.Recv(msg)
+	if err != io.EOF {
+		messageType, version := getMessageType(msg)
+		newrelic.GetSecurityAgentInterface().SendEvent("GRPC", msg, messageType, version)
+	}
+	return err
+}
+func (s wrappedServerStream) Error() error {
+	return s.stream.Error()
+}
+func (s wrappedServerStream) Close() error {
+	return s.stream.Close()
+}
+
+func getMessageType(req any) (string, string) {
+	messageType := ""
+	version := "v2"
+	messagev2, ok := req.(protoV2.Message)
+	if ok {
+		messageType = string(messagev2.ProtoReflect().Descriptor().FullName())
+	} else {
+		messagev1, ok := req.(protoV1.Message)
+		if ok {
+			messageType = string(protoV1.MessageReflect(messagev1).Descriptor().FullName())
+			version = "v1"
+		}
+	}
+	return messageType, version
+}

v3/integrations/nrsecurityagent/README.md

@@ -54,6 +54,8 @@ validator_service_url: wss://csec.nr-data.net
 detection:
   rxss:
     enabled: true
+request:
+  body_limit:1
 ```
 
 * Based on additional packages imported by the user application, add suitable instrumentation package imports. 

v3/integrations/nrsecurityagent/go.mod

@@ -3,7 +3,7 @@ module github.com/newrelic/go-agent/v3/integrations/nrsecurityagent
 go 1.19
 
 require (
-	github.com/newrelic/csec-go-agent v0.4.0
+	github.com/newrelic/csec-go-agent v0.5.1
 	github.com/newrelic/go-agent/v3 v3.26.0
 	github.com/newrelic/go-agent/v3/integrations/nrsqlite3 v1.2.0
 	gopkg.in/yaml.v2 v2.4.0

v3/integrations/nrsecurityagent/nrsecurityagent.go

@@ -30,6 +30,7 @@ func defaultSecurityConfig() SecurityConfig {
 	cfg.Security.Mode = "IAST"
 	cfg.Security.Agent.Enabled = true
 	cfg.Security.Detection.Rxss.Enabled = true
+	cfg.Security.Request.BodyLimit = 300
 	return cfg
 }
 
@@ -108,6 +109,8 @@ func ConfigSecurityFromYaml() ConfigOption {
 //	NEW_RELIC_SECURITY_MODE						scanning mode: "IAST" for now
 //	NEW_RELIC_SECURITY_AGENT_ENABLED			(boolean)
 //	NEW_RELIC_SECURITY_DETECTION_RXSS_ENABLED	(boolean)
+//	NEW_RELIC_SECURITY_REQUEST_BODY_LIMIT		(integer) set limit on read request body in kb. By default, this is "300"
+
 func ConfigSecurityFromEnvironment() ConfigOption {
 	return func(cfg *SecurityConfig) {
 		assignBool := func(field *bool, name string) {
@@ -125,11 +128,22 @@ func ConfigSecurityFromEnvironment() ConfigOption {
 			}
 		}
 
+		assignInt := func(field *int, name string) {
+			if env := os.Getenv(name); env != "" {
+				if i, err := strconv.Atoi(env); nil != err {
+					cfg.Error = fmt.Errorf("invalid %s value: %s", name, env)
+				} else {
+					*field = i
+				}
+			}
+		}
+
 		assignBool(&cfg.Security.Enabled, "NEW_RELIC_SECURITY_ENABLED")
 		assignString(&cfg.Security.Validator_service_url, "NEW_RELIC_SECURITY_VALIDATOR_SERVICE_URL")
 		assignString(&cfg.Security.Mode, "NEW_RELIC_SECURITY_MODE")
 		assignBool(&cfg.Security.Agent.Enabled, "NEW_RELIC_SECURITY_AGENT_ENABLED")
 		assignBool(&cfg.Security.Detection.Rxss.Enabled, "NEW_RELIC_SECURITY_DETECTION_RXSS_ENABLED")
+		assignInt(&cfg.Security.Request.BodyLimit, "NEW_RELIC_SECURITY_REQUEST_BODY_LIMIT")
 	}
 }
 
@@ -160,3 +174,10 @@ func ConfigSecurityEnable(isEnabled bool) ConfigOption {
 		cfg.Security.Enabled = isEnabled
 	}
 }
+
+// ConfigSecurityRequestBodyLimit set limit on read request body in kb. By default, this is "300"
+func ConfigSecurityRequestBodyLimit(bodyLimit int) ConfigOption {
+	return func(cfg *SecurityConfig) {
+		cfg.Security.Request.BodyLimit = bodyLimit
+	}
+}

v3/newrelic/secure_agent.go

@@ -4,18 +4,15 @@ import (
 	"net/http"
 )
 
-//
 // secureAgent is a global interface point for the nrsecureagent's hooks into the go agent.
 // The default value for this is a noOpSecurityAgent value, which has null definitions for
 // the methods. The Go compiler is expected to optimize away all the securityAgent method
 // calls in this case, effectively removing the hooks from the running agent.
 //
 // If the nrsecureagent integration was initialized, it will register a real securityAgent
 // value in the securityAgent varialble instead, thus "activating" the hooks.
-//
 var secureAgent securityAgent = noOpSecurityAgent{}
 
-//
 // GetSecurityAgentInterface returns the securityAgent value
 // which provides the working interface to the installed
 // security agent (or to a no-op interface if none were
@@ -26,7 +23,6 @@ var secureAgent securityAgent = noOpSecurityAgent{}
 // This avoids exposing the variable itself so it's not
 // writable externally and also sets up for the future if this
 // ends up not being a global variable later.
-//
 func GetSecurityAgentInterface() securityAgent {
 	return secureAgent
 }
@@ -38,6 +34,7 @@ type securityAgent interface {
 	IsSecurityActive() bool
 	DistributedTraceHeaders(hdrs *http.Request, secureAgentevent any)
 	SendExitEvent(any, error)
+	RequestBodyReadLimit() int
 }
 
 func (app *Application) RegisterSecurityAgent(s securityAgent) {
@@ -88,13 +85,62 @@ func (t noOpSecurityAgent) DistributedTraceHeaders(hdrs *http.Request, secureAge
 
 func (t noOpSecurityAgent) SendExitEvent(secureAgentevent any, err error) {
 }
+func (t noOpSecurityAgent) RequestBodyReadLimit() int {
+	return 300 * 1000
+}
 
-//
 // IsSecurityAgentPresent returns true if there's an actual security agent hooked in to the
 // Go APM agent, whether or not it's enabled or operating in any particular mode. It returns
 // false only if the hook-in interface for those functions is a No-Op will null functionality.
-//
 func IsSecurityAgentPresent() bool {
 	_, isNoOp := secureAgent.(noOpSecurityAgent)
 	return !isNoOp
 }
+
+type BodyBuffer struct {
+	buf             []byte
+	isDataTruncated bool
+}
+
+func (b *BodyBuffer) Write(p []byte) (int, error) {
+	if l := len(b.buf); len(p) <= secureAgent.RequestBodyReadLimit()-l {
+		b.buf = append(b.buf, p...)
+		return len(p), nil
+	} else if l := len(b.buf); secureAgent.RequestBodyReadLimit()-l > 1 {
+		end := secureAgent.RequestBodyReadLimit() - l
+		b.buf = append(b.buf, p[:end-1]...)
+		return end, nil
+	} else {
+		b.isDataTruncated = true
+		return 0, nil
+	}
+}
+
+func (b *BodyBuffer) Len() int {
+	if b == nil {
+		return 0
+	}
+	return len(b.buf)
+
+}
+
+func (b *BodyBuffer) read() []byte {
+	if b == nil {
+		return make([]byte, 0)
+	}
+	return b.buf
+}
+
+func (b *BodyBuffer) isBodyTruncated() bool {
+	if b == nil {
+		return false
+	}
+	return b.isDataTruncated
+}
+func (b *BodyBuffer) String() (string, bool) {
+	if b == nil {
+		return "", false
+	}
+	return string(b.buf), b.isDataTruncated
+
+}

v3/newrelic/transaction.go

@@ -4,7 +4,6 @@
 package newrelic
 
 import (
-	"bytes"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -246,20 +245,14 @@ func serverName(r *http.Request) string {
 	return ""
 }
 
-func reqBody(req *http.Request) []byte {
-	var bodyBuffer bytes.Buffer
-	requestBuffer := make([]byte, 0)
-	bodyReader := io.TeeReader(req.Body, &bodyBuffer)
-
-	if bodyReader != nil && req.Body != nil {
-		reqBuffer, err := io.ReadAll(bodyReader)
-		if err == nil {
-			requestBuffer = reqBuffer
-		}
-		r := io.NopCloser(bytes.NewBuffer(requestBuffer))
-		req.Body = r
+func reqBody(req *http.Request) *BodyBuffer {
+	if IsSecurityAgentPresent() {
+		buf := &BodyBuffer{buf: make([]byte, 0, 100)}
+		tee := io.TeeReader(req.Body, buf)
+		req.Body = io.NopCloser(tee)
+		return buf
 	}
-	return bytes.TrimRight(requestBuffer, "\x00")
+	return nil
 }
 
 // SetWebRequest marks the transaction as a web transaction.  SetWebRequest
@@ -607,7 +600,7 @@ type WebRequest struct {
 
 	// The following fields are needed for the secure agent's vulnerability
 	// detection features.
-	Body          []byte
+	Body          *BodyBuffer
 	ServerName    string
 	Type          string
 	RemoteAddress string
@@ -634,7 +627,17 @@ func (webrequest WebRequest) GetHost() string {
 }
 
 func (webrequest WebRequest) GetBody() []byte {
-	return webrequest.Body
+	if webrequest.Body == nil {
+		return make([]byte, 0)
+	}
+	return webrequest.Body.read()
+}
+
+func (webrequest WebRequest) IsDataTruncated() bool {
+	if webrequest.Body == nil {
+		return false
+	}
+	return webrequest.Body.isBodyTruncated()
 }
 
 func (webrequest WebRequest) GetServerName() string {