corerequestcontext.go

package common

import (
	"context"
	"net/http"

	"github.com/anz-bank/sysl-go/logconfig"

	"github.com/anz-bank/sysl-go/common/internal"

	"github.com/anz-bank/pkg/log"
	"github.com/sirupsen/logrus"
)

// Deprecated: Use ServerParams.WithPkgLogger instead
func GetLogEntryFromContext(ctx context.Context) *logrus.Entry {
	core := getCoreContext(ctx)
	if core == nil {
		return nil
	}
	return core.entry
}

// Deprecated: Use ServerParams.WithPkgLogger instead
func GetLoggerFromContext(ctx context.Context) *logrus.Logger {
	core := getCoreContext(ctx)
	if core == nil {
		return nil
	}
	return core.logger
}

func NewLoggingRoundTripper(name string, base http.RoundTripper) http.RoundTripper {
	// temporary pass-through to get the real round tripper from the request context
	return &tempRoundtripper{name, base}
}

type coreRequestContext struct {
	logger *logrus.Logger
	entry  *logrus.Entry
}

type reqHeaderContext struct {
	header http.Header
}

type respHeaderAndStatusContext struct {
	header http.Header
	status int
}

// LoggerToContext create a new context containing the logger
// Deprecated: Use ServerParams.WithPkgLogger instead
func LoggerToContext(ctx context.Context, logger *logrus.Logger, entry *logrus.Entry) context.Context {
	return context.WithValue(ctx, coreRequestContextKey{}, &coreRequestContext{logger, entry})
}

// RequestHeaderToContext create a new context containing the request header
func RequestHeaderToContext(ctx context.Context, header http.Header) context.Context {
	return context.WithValue(ctx, reqHeaderContextKey{}, &reqHeaderContext{header})
}

// RequestHeaderFromContext retrieve the request header from the context
func RequestHeaderFromContext(ctx context.Context) http.Header {
	reqHeader := getReqHeaderContext(ctx)

	if reqHeader == nil {
		return nil
	}
	return reqHeader.header
}

// RespHeaderAndStatusToContext create a new context containing the response header and status
func RespHeaderAndStatusToContext(ctx context.Context, header http.Header, status int) context.Context {
	return context.WithValue(ctx, respHeaderAndStatusContextKey{}, &respHeaderAndStatusContext{header, status})
}

// RespHeaderAndStatusFromContext retrieve response header and status from the context
func RespHeaderAndStatusFromContext(ctx context.Context) (header http.Header, status int) {
	respHeaderAndStatus := getRespHeaderAndStatusContext(ctx)

	if respHeaderAndStatus == nil {
		return nil, http.StatusOK
	}

	header = respHeaderAndStatus.header
	status = respHeaderAndStatus.status
	return
}

func UpdateResponseStatus(ctx context.Context, status int) error {
	respHeaderAndStatus := getRespHeaderAndStatusContext(ctx)

	if respHeaderAndStatus == nil {
		return CreateError(ctx, InternalError, "response status not in context", nil)
	}
	respHeaderAndStatus.status = status
	return nil
}

func CoreRequestContextMiddleware(logger *logrus.Logger) func(next http.Handler) http.Handler {
	ctx := LoggerToContext(context.Background(), logger, nil)
	return CoreRequestContextMiddlewareWithContext(ctx)
}

func CoreRequestContextMiddlewareWithContext(ctx context.Context) func(next http.Handler) http.Handler {
	ctxlogger := GetLoggerFromContext(ctx)
	pkgLogger := log.From(ctx)
	verboseLogging := logconfig.IsVerboseLogging(ctx)
	return func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			ctx := r.Context()
			ctx = LoggerToContext(ctx, ctxlogger, nil)
			ctx = log.WithLogger(pkgLogger).Onto(ctx)
			ctx = log.With(traceIDLogField, GetTraceIDFromContext(ctx)).Onto(ctx)
			ctx = logconfig.SetVerboseLogging(ctx, verboseLogging)

			ctx = internal.AddResponseBodyMonitorToContext(ctx)
			defer internal.CheckForUnclosedResponses(ctx)
			reqLogger, entry := internal.NewRequestLogger(ctx, r)
			w = reqLogger.ResponseWriter(w)
			defer reqLogger.FlushLog()

			r = r.WithContext(ctx)

			tl := internal.NewRequestTimer(w, r)
			w = tl.RespWrapper
			defer tl.Log(entry)

			next.ServeHTTP(w, r)
		}
		return http.HandlerFunc(fn)
	}
}

type coreRequestContextKey struct{}

func getCoreContext(ctx context.Context) *coreRequestContext {
	coreRequestCtx := ctx.Value(coreRequestContextKey{})
	if coreRequestCtx == nil {
		return nil
	}
	return coreRequestCtx.(*coreRequestContext)
}

type reqHeaderContextKey struct{}
type respHeaderAndStatusContextKey struct{}

func getReqHeaderContext(ctx context.Context) *reqHeaderContext {
	reqHeaderCtx := ctx.Value(reqHeaderContextKey{})
	if reqHeaderCtx == nil {
		return nil
	}
	return reqHeaderCtx.(*reqHeaderContext)
}

func getRespHeaderAndStatusContext(ctx context.Context) *respHeaderAndStatusContext {
	respHeaderAndStatusCtx := ctx.Value(respHeaderAndStatusContextKey{})
	if respHeaderAndStatusCtx == nil {
		return nil
	}
	return respHeaderAndStatusCtx.(*respHeaderAndStatusContext)
}

type tempRoundtripper struct {
	name string
	base http.RoundTripper
}

func (t *tempRoundtripper) RoundTrip(r *http.Request) (*http.Response, error) {
	ctx := log.With("Downsteam", t.name).Onto(r.Context())
	return internal.NewLoggingRoundTripper(ctx, t.base).RoundTrip(r)
}