Skip to content

Commit

Permalink
internal/appsec: refactor blocking actions to call Block()
Browse files Browse the repository at this point in the history
Signed-off-by: Eliott Bouhana <eliott.bouhana@datadoghq.com>
  • Loading branch information
eliottness committed Nov 14, 2024
1 parent 219afa9 commit be7d55b
Show file tree
Hide file tree
Showing 6 changed files with 192 additions and 99 deletions.
58 changes: 19 additions & 39 deletions internal/appsec/emitter/httpsec/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -56,20 +56,31 @@ type (
func (HandlerOperationArgs) IsArgOf(*HandlerOperation) {}
func (HandlerOperationRes) IsResultOf(*HandlerOperation) {}

func StartOperation(ctx context.Context, args HandlerOperationArgs) (*HandlerOperation, *atomic.Pointer[actions.BlockHTTP], context.Context) {
wafOp, ctx := waf.StartContextOperation(ctx)
func StartOperation(w http.ResponseWriter, r *http.Request, pathParams map[string]string, opts *Config) (*HandlerOperation, *atomic.Bool, context.Context) {
wafOp, ctx := waf.StartContextOperation(r.Context())
op := &HandlerOperation{
Operation: dyngo.NewOperation(wafOp),
ContextOperation: wafOp,
}

// We need to use an atomic pointer to store the action because the action may be created asynchronously in the future
var action atomic.Pointer[actions.BlockHTTP]
var blocked atomic.Bool
dyngo.OnData(op, func(a *actions.BlockHTTP) {
action.Store(a)
a.Handler.ServeHTTP(w, r)
for _, f := range opts.OnBlock {
f()
}
})

return op, &action, dyngo.StartAndRegisterOperation(ctx, op, args)
return op, &blocked, dyngo.StartAndRegisterOperation(ctx, op, HandlerOperationArgs{
Method: r.Method,
RequestURI: r.RequestURI,
Host: r.Host,
RemoteAddr: r.RemoteAddr,
Headers: r.Header,
Cookies: makeCookies(r.Cookies()),
QueryParams: r.URL.Query(),
PathParams: pathParams,
})
}

// Finish the HTTP handler operation and its children operations and write everything to the service entry span.
Expand Down Expand Up @@ -125,18 +136,7 @@ func BeforeHandle(
opts.ResponseHeaderCopier = defaultWrapHandlerConfig.ResponseHeaderCopier
}

op, blockAtomic, ctx := StartOperation(r.Context(), HandlerOperationArgs{
Method: r.Method,
RequestURI: r.RequestURI,
Host: r.Host,
RemoteAddr: r.RemoteAddr,
Headers: r.Header,
Cookies: makeCookies(r.Cookies()),
QueryParams: r.URL.Query(),
PathParams: pathParams,
})
tr := r.WithContext(ctx)
var blocked atomic.Bool
op, blocked, ctx := StartOperation(w, r, pathParams, opts)

afterHandle := func() {
var statusCode int
Expand All @@ -147,28 +147,9 @@ func BeforeHandle(
Headers: opts.ResponseHeaderCopier(w),
StatusCode: statusCode,
}, span)

if blockPtr := blockAtomic.Swap(nil); blockPtr != nil {
blockPtr.Handler.ServeHTTP(w, tr)
blocked.Store(true)
}

// Execute the onBlock functions to make sure blocking works properly
// in case we are instrumenting the Gin framework
if blocked.Load() {
for _, f := range opts.OnBlock {
f()
}
}
}

if blockPtr := blockAtomic.Swap(nil); blockPtr != nil {
// handler is replaced
blockPtr.Handler.ServeHTTP(w, tr)
blocked.Store(true)
}

return w, tr, afterHandle, blocked.Load()
return w, r.WithContext(ctx), afterHandle, blocked.Load()
}

// WrapHandler wraps the given HTTP handler with the abstract HTTP operation defined by HandlerOperationArgs and
Expand All @@ -177,7 +158,6 @@ func BeforeHandle(
// It is a specific patch meant for Gin, for which we must abort the
// context since it uses a queue of handlers and it's the only way to make
// sure other queued handlers don't get executed.
// TODO: this patch must be removed/improved when we rework our actions/operations system
func WrapHandler(handler http.Handler, span ddtrace.Span, pathParams map[string]string, opts *Config) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
tw, tr, afterHandle, handled := BeforeHandle(w, r, span, pathParams, opts)
Expand Down
2 changes: 1 addition & 1 deletion internal/appsec/emitter/waf/actions/actions.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ var actionHandlers = map[string]actionHandler{}

func registerActionHandler(aType string, handler actionHandler) {
if _, ok := actionHandlers[aType]; ok {
log.Warn("appsec: action type `%s` already registered", aType)
log.Debug("appsec: action type `%s` already registered", aType)
return
}
actionHandlers[aType] = handler
Expand Down
6 changes: 3 additions & 3 deletions internal/appsec/emitter/waf/actions/actions_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@ import (
func TestNewHTTPBlockRequestAction(t *testing.T) {
mux := http.NewServeMux()
srv := httptest.NewServer(mux)
mux.HandleFunc("/json", newHTTPBlockRequestAction(403, "json").ServeHTTP)
mux.HandleFunc("/html", newHTTPBlockRequestAction(403, "html").ServeHTTP)
mux.HandleFunc("/auto", newHTTPBlockRequestAction(403, "auto").ServeHTTP)
mux.HandleFunc("/json", newHTTPBlockRequestAction(403, BlockingTemplateJSON).ServeHTTP)
mux.HandleFunc("/html", newHTTPBlockRequestAction(403, BlockingTemplateHTML).ServeHTTP)
mux.HandleFunc("/auto", newHTTPBlockRequestAction(403, BlockingTemplateAuto).ServeHTTP)
defer srv.Close()

t.Run("json", func(t *testing.T) {
Expand Down
144 changes: 100 additions & 44 deletions internal/appsec/emitter/waf/actions/block.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,15 +47,23 @@ func init() {
registerActionHandler("block_request", NewBlockAction)
}

const (
BlockingTemplateJSON blockingTemplateType = "json"
BlockingTemplateHTML blockingTemplateType = "html"
BlockingTemplateAuto blockingTemplateType = "auto"
)

type (
blockingTemplateType string

// blockActionParams are the dynamic parameters to be provided to a "block_request"
// action type upon invocation
blockActionParams struct {
// GRPCStatusCode is the gRPC status code to be returned. Since 0 is the OK status, the value is nullable to
// be able to distinguish between unset and defaulting to Abort (10), or set to OK (0).
GRPCStatusCode *int `mapstructure:"grpc_status_code,omitempty"`
StatusCode int `mapstructure:"status_code"`
Type string `mapstructure:"type,omitempty"`
GRPCStatusCode *int `mapstructure:"grpc_status_code,omitempty"`
StatusCode int `mapstructure:"status_code"`
Type blockingTemplateType `mapstructure:"type,omitempty"`
}
// GRPCWrapper is an opaque prototype abstraction for a gRPC handler (to avoid importing grpc)
// that returns a status code and an error
Expand All @@ -70,6 +78,12 @@ type (
BlockHTTP struct {
http.Handler
}

HTTPBlockHandlerConfig struct {
Template []byte
ContentType string
StatusCode int
}
)

func (a *BlockGRPC) EmitData(op dyngo.Operation) {
Expand All @@ -83,32 +97,28 @@ func (a *BlockHTTP) EmitData(op dyngo.Operation) {
}

func newGRPCBlockRequestAction(status int) *BlockGRPC {
return &BlockGRPC{GRPCWrapper: newGRPCBlockHandler(status)}
}

func newGRPCBlockHandler(status int) GRPCWrapper {
return func() (uint32, error) {
return &BlockGRPC{GRPCWrapper: func() (uint32, error) {
return uint32(status), &events.BlockingSecurityEvent{}
}
}}
}

func blockParamsFromMap(params map[string]any) (blockActionParams, error) {
grpcCode := 10
p := blockActionParams{
Type: "auto",
parsedParams := blockActionParams{
Type: BlockingTemplateAuto,
StatusCode: 403,
GRPCStatusCode: &grpcCode,
}

if err := mapstructure.WeakDecode(params, &p); err != nil {
return p, err
if err := mapstructure.WeakDecode(params, &parsedParams); err != nil {
return parsedParams, err
}

if p.GRPCStatusCode == nil {
p.GRPCStatusCode = &grpcCode
if parsedParams.GRPCStatusCode == nil {
parsedParams.GRPCStatusCode = &grpcCode
}

return p, nil
return parsedParams, nil
}

// NewBlockAction creates an action for the "block_request" action type
Expand All @@ -124,38 +134,84 @@ func NewBlockAction(params map[string]any) []Action {
}
}

func newHTTPBlockRequestAction(status int, template string) *BlockHTTP {
return &BlockHTTP{Handler: newBlockHandler(status, template)}
func newHTTPBlockRequestAction(statusCode int, template blockingTemplateType) *BlockHTTP {
return &BlockHTTP{Handler: http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
template := template
if template == BlockingTemplateAuto {
template = blockingTemplateTypeFromHeaders(request.Header)
}

if UnwrapGetStatusCode(writer) != 0 {
// The status code has already been set, so we can't change it, do nothing
return
}

blocker, found := UnwrapBlocker(writer)
if found {
// We found our custom response writer, so we can block futur calls to Write and WriteHeader
defer blocker()
}

writer.Header().Set("Content-Type", template.ContentType())
writer.WriteHeader(statusCode)
writer.Write(template.Template())
})}
}

// newBlockHandler creates, initializes and returns a new BlockRequestAction
func newBlockHandler(status int, template string) http.Handler {
htmlHandler := newBlockRequestHandler(status, "text/html", blockedTemplateHTML)
jsonHandler := newBlockRequestHandler(status, "application/json", blockedTemplateJSON)
switch template {
case "json":
return jsonHandler
case "html":
return htmlHandler
default:
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
h := jsonHandler
hdr := r.Header.Get("Accept")
htmlIdx := strings.Index(hdr, "text/html")
jsonIdx := strings.Index(hdr, "application/json")
// Switch to html handler if text/html comes before application/json in the Accept header
if htmlIdx != -1 && (jsonIdx == -1 || htmlIdx < jsonIdx) {
h = htmlHandler
}
h.ServeHTTP(w, r)
})
func blockingTemplateTypeFromHeaders(headers http.Header) blockingTemplateType {
hdr := headers.Get("Accept")
htmlIdx := strings.Index(hdr, "text/html")
jsonIdx := strings.Index(hdr, "application/json")
// Switch to html handler if text/html comes before application/json in the Accept header
if htmlIdx != -1 && (jsonIdx == -1 || htmlIdx < jsonIdx) {
return BlockingTemplateHTML
}

return BlockingTemplateJSON
}

func (typ blockingTemplateType) Template() []byte {
if typ == BlockingTemplateHTML {
return blockedTemplateHTML
}

return blockedTemplateJSON
}

func (typ blockingTemplateType) ContentType() string {
if typ == BlockingTemplateHTML {
return "text/html"
}

return "application/json"
}

// UnwrapBlocker unwraps the right struct method from contrib/internal/httptrace.responseWriter
// and returns the Block() function and if it was found.
func UnwrapBlocker(writer http.ResponseWriter) (func(), bool) {
// this is part of the contrib/internal/httptrace.responseWriter interface
wrapped, ok := writer.(interface {
Block()
})
if !ok {
// Somehow we can't access the wrapped response writer, so we can't block the response
return nil, false
}

return wrapped.Block, true
}

func newBlockRequestHandler(status int, ct string, payload []byte) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", ct)
w.WriteHeader(status)
w.Write(payload)
// UnwrapGetStatusCode unwraps the right struct method from contrib/internal/httptrace.responseWriter
// and calls it to know if a call to WriteHeader has been made and returns the status code.
func UnwrapGetStatusCode(writer http.ResponseWriter) int {
// this is part of the contrib/internal/httptrace.responseWriter interface
wrapped, ok := writer.(interface {
Status() int
})
if !ok {
// Somehow we can't access the wrapped response writer, so we can't get the status code
return 0
}

return wrapped.Status()
}
25 changes: 20 additions & 5 deletions internal/appsec/emitter/waf/actions/http_redirect.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,9 @@ func init() {
}

func redirectParamsFromMap(params map[string]any) (redirectActionParams, error) {
var p redirectActionParams
err := mapstructure.WeakDecode(params, &p)
return p, err
var parsedParams redirectActionParams
err := mapstructure.WeakDecode(params, &parsedParams)
return parsedParams, err
}

func newRedirectRequestAction(status int, loc string) *BlockHTTP {
Expand All @@ -38,9 +38,24 @@ func newRedirectRequestAction(status int, loc string) *BlockHTTP {

// If location is not set we fall back on a default block action
if loc == "" {
return &BlockHTTP{Handler: newBlockHandler(http.StatusForbidden, string(blockedTemplateJSON))}
return newHTTPBlockRequestAction(http.StatusForbidden, BlockingTemplateAuto)
}
return &BlockHTTP{Handler: http.RedirectHandler(loc, status)}

redirectHandler := http.RedirectHandler(loc, status)
return &BlockHTTP{Handler: http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if UnwrapGetStatusCode(writer) != 0 {
// The status code has already been set, so we can't change it, do nothing
return
}

blocker, found := UnwrapBlocker(writer)
if found {
// We found our custom response writer, so we can block futur calls to Write and WriteHeader
defer blocker()
}

redirectHandler.ServeHTTP(writer, request)
})}
}

// NewRedirectAction creates an action for the "redirect_request" action type
Expand Down
Loading

0 comments on commit be7d55b

Please sign in to comment.