middleware.go

package httpsling

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"os"

	"github.com/theopenlane/utils/rout"
)

// Middleware can be used to wrap Doers with additional functionality.
type Middleware func(Doer) Doer

// Apply implements Option
func (m Middleware) Apply(r *Requester) error {
	r.Middleware = append(r.Middleware, m)

	return nil
}

// Wrap applies a set of middleware to a Doer.  The returned Doer will invoke
// the middleware in the order of the arguments.
func Wrap(d Doer, m ...Middleware) Doer {
	for i := len(m) - 1; i > -1; i-- {
		d = m[i](d)
	}

	return d
}

// Dump dumps requests and responses to a writer.  Just intended for debugging.
func Dump(w io.Writer) Middleware {
	return func(next Doer) Doer {
		return DoerFunc(func(req *http.Request) (*http.Response, error) {
			dump, dumperr := httputil.DumpRequestOut(req, true)
			if dumperr != nil {
				io.WriteString(w, "Error dumping request: "+dumperr.Error()+"\n") // nolint: errcheck
			} else {
				io.WriteString(w, string(dump)+"\n") // nolint: errcheck
			}

			resp, err := next.Do(req)
			if resp != nil {
				dump, dumperr = httputil.DumpResponse(resp, true)
				if dumperr != nil {
					io.WriteString(w, "Error dumping response: "+dumperr.Error()+"\n") // nolint: errcheck
				} else {
					io.WriteString(w, string(dump)+"\n") // nolint: errcheck
				}
			}

			return resp, err
		})
	}
}

// DumpToStout dumps requests and responses to os.Stdout
func DumpToStout() Middleware {
	return Dump(os.Stdout)
}

// DumpToStderr dumps requests and responses to os.Stderr
func DumpToStderr() Middleware {
	return Dump(os.Stderr)
}

type logFunc func(a ...interface{})

func (f logFunc) Write(p []byte) (n int, err error) {
	f(string(p))
	return len(p), nil
}

// DumpToLog dumps the request and response to a logging function.
// logf is compatible with fmt.Print(), testing.T.Log, or log.XXX()
// functions.
//
// logf will be invoked once for the request, and once for the response.
// Each invocation will only have a single argument (the entire request
// or response is logged as a single string value).
func DumpToLog(logf func(a ...interface{})) Middleware {
	return Dump(logFunc(logf))
}

// ExpectCode generates an error if the response's status code does not match
// the expected code.
//
// The response body will still be read and returned.
func ExpectCode(code int) Middleware {
	return func(next Doer) Doer {
		return DoerFunc(func(req *http.Request) (*http.Response, error) {
			r, c := getCodeChecker(req)
			c.code = code
			resp, err := next.Do(r)

			return c.checkCode(resp, err)
		})
	}
}

// ExpectSuccessCode is middleware which generates an error if the response's status code is not between 200 and
// 299.
//
// The response body will still be read and returned.
func ExpectSuccessCode() Middleware {
	return func(next Doer) Doer {
		return DoerFunc(func(req *http.Request) (*http.Response, error) {
			r, c := getCodeChecker(req)
			c.code = expectSuccessCode
			resp, err := next.Do(r)

			return c.checkCode(resp, err)
		})
	}
}

type ctxKey int

const expectCodeCtxKey ctxKey = iota

const expectSuccessCode = -1

type codeChecker struct {
	code int
}

func (c *codeChecker) checkCode(resp *http.Response, err error) (*http.Response, error) {
	switch {
	case err != nil, resp == nil:
	case c.code == expectSuccessCode:
		if resp.StatusCode < 200 || resp.StatusCode > 299 {
			err = rout.HTTPErrorResponse(
				fmt.Errorf("%w: server returned unsuccessful status code: %d",
					ErrUnsuccessfulResponse,
					resp.StatusCode,
				))
		}
	case c.code != resp.StatusCode:
		err = rout.HTTPErrorResponse(
			fmt.Errorf("%w: server returned unexpected status code. expected: %d, received: %d",
				ErrUnsuccessfulResponse,
				c.code,
				resp.StatusCode,
			))
	}

	return resp, err
}

func getCodeChecker(req *http.Request) (*http.Request, *codeChecker) {
	c, _ := req.Context().Value(expectCodeCtxKey).(*codeChecker)
	if c == nil {
		c = &codeChecker{}
		req = req.WithContext(context.WithValue(req.Context(), expectCodeCtxKey, c))
	}

	return req, c
}